"""Orchestrator Agent — Steps.

DAG:
  classify_query | guardian_validate  (parallel)
    -> gate_decision
      -> call_planner
        -> extract_plan
          -> build_workflow
"""

import json

from app import agent
from abi_core.common.utils import abi_logging
from abi_core.common.a2a_response import A2AResponse
from abi_core.common.semantic_tools import tool_find_agent, MCPToolkit
from abi_core.common.workflow import AgentInteractionFlow, InteractionFlowNode
from abi_core.common import prompts
from a2a.types import AgentCard
from abi_core.common.agent_card_loader import build_agent_card, get_agent_url
from abi_core.agent.agent import AbiAgent
from config import AGENT_CARD, config

# Agents that must NEVER be deregistered (infrastructure)
INFRA_AGENTS = {"builder", "planner", "orchestrator", "guardian", "semantic-layer"}

# Working-memory topic used to mark a session awaiting a planner clarification
PENDING_CLARIFICATION_TOPIC = "pending_clarification"


@agent.step(
    name="classify_query",
    input_map={"query": "$input.query", "context_id": "$input.context_id"},
)
async def classify_query(query, context_id=""):
    """Classify query as simple or complex using the orchestrator's LLM.

    If the session has a pending clarification (a system-level event stored in
    short-term memory), the user's message is the answer to that clarification.
    In that case we skip the probabilistic LLM triage and deterministically
    route back to the planner with the original query + the answer.
    """
    from abi_core.agent.llm_provider import invoke
    from abi_core.common.utils import clean_llm_json

    # ── Deterministic check: is this a pending clarification answer? ──
    pending = await _load_pending_clarification(context_id)
    if pending:
        abi_logging(f"[🔁] Clarification answer detected for session '{context_id}'")
        return {
            "classification": "clarification_response",
            "original_query": pending.get("original_query", ""),
            "clarification": pending.get("clarification", ""),
            "user_answer": query,
        }

    try:
        text = await invoke(
            config.LLM_CONFIG,
            prompts.ORCHESTRATOR_TRIAGE_PROMPT.format(query=query),
        )
        parsed = clean_llm_json(text)
        classification = parsed.get("classification", "complex")

        if classification not in ("simple", "complex"):
            classification = "complex"

        abi_logging(f"[🔍] Triage: '{classification}' for query: {query[:80]}")
        return {"classification": classification}

    except Exception as e:
        abi_logging(f"[⚠️] Triage failed, defaulting to complex: {e}")
        return {"classification": "complex"}


async def _load_pending_clarification(context_id: str) -> dict:
    """Return the pending-clarification event for a session, or {} if none.

    The event is a system-level signal stored in short-term (working) memory by
    the orchestrator when it emits an ``input_required``. Reading it is best
    effort: if memory is unavailable, we degrade to normal triage (returns {}).

    On a hit, the event is cleared so the next normal request in the session is
    not misinterpreted as another clarification answer.
    """
    if not config.AGENT_MEMORY_URL or not context_id:
        return {}

    try:
        from abi_core.memory import get_short_term_memory

        raw = await get_short_term_memory(
            context_id=context_id, memory_url=config.AGENT_MEMORY_URL
        )
        if not raw:
            return {}

        # Working memory accumulates; a clarification is ACTIVE only if the most
        # recent clarification-related marker is a 'pending' one (not 'resolved').
        last_pending = None
        for line in raw.splitlines():
            if "clarification_resolved" in line:
                last_pending = None  # a later resolution cancels an earlier pending
            elif PENDING_CLARIFICATION_TOPIC in line:
                last_pending = line
        if not last_pending:
            return {}

        payload = _extract_clarification_payload(last_pending)
        if payload:
            await _clear_pending_clarification(context_id)
        return payload
    except Exception as e:
        abi_logging(f"[⚠️] Could not read pending clarification: {e}")
        return {}


def _extract_clarification_payload(text: str) -> dict:
    """Extract the JSON payload embedded in a pending_clarification memory line."""
    import json
    start = text.find("{")
    end = text.rfind("}")
    if start == -1 or end == -1 or end <= start:
        return {}
    try:
        return json.loads(text[start:end + 1])
    except (json.JSONDecodeError, ValueError):
        return {}


async def _clear_pending_clarification(context_id: str) -> None:
    """Mark the pending clarification as resolved in short-term memory."""
    try:
        from abi_core.memory import add_short_term_memory

        await add_short_term_memory(
            topic="clarification_resolved",
            task=context_id,
            content="clarification_resolved marker",
            context_id=context_id,
            memory_url=config.AGENT_MEMORY_URL,
        )
    except Exception as e:
        abi_logging(f"[⚠️] Could not clear pending clarification: {e}")


@agent.step(
    name="guardian_validate",
    input_map={
        "query": "$input.query",
        "context_id": "$input.context_id",
    },
)
async def guardian_validate(query, context_id):
    """Call Guardian agent to validate query security.

    Checks: prompt injection, reverse engineering, policy compliance.
    Returns dict with 'allowed', 'reason', and 'status'.
    """
    try:
        guardian_card = await tool_find_agent.ainvoke({"query": "guardian"})
        if not guardian_card:
            abi_logging("[⚠️] Guardian agent not found — cannot validate")
            return {
                "status": "error",
                "allowed": False,
                "reason": "Guardian service unavailable",
            }

        # Call guardian via A2A
        validation_query = json.dumps({
            "action": "validate_query",
            "query": query,
            "context_id": context_id,
            "checks": ["prompt_injection", "reverse_engineering", "policy_compliance"],
        })

        workflow = AgentInteractionFlow()
        node = InteractionFlowNode(
            task=validation_query,
            source_agent_card=AGENT_CARD,
            target_agent_card=guardian_card,
            node_key="guardian_gate",
            node_label="Security Validation",
        )
        workflow.add_node(node)
        workflow.set_node_attributes(
            node.id, {"context_id": context_id, "query": validation_query}
        )
        workflow.set_source_card(AGENT_CARD)

        results = []
        async for chunk in workflow.run_workflow():
            results.append(chunk)

        # Parse guardian response
        for resp in A2AResponse.from_results(results):
            if resp.data:
                allowed = resp.data.get("allowed", True)
                return {
                    "status": "blocked" if not allowed else "approved",
                    "allowed": allowed,
                    "reason": resp.data.get("reason", ""),
                    "risk_score": resp.data.get("risk_score", 0.0),
                }
            if resp.text:
                # Guardian responded with text — assume approved
                abi_logging(f"[🛡️] Guardian text response: {resp.text[:100]}")
                return {"status": "approved", "allowed": True, "reason": resp.text}

        # No parseable response — treat as approved with warning
        abi_logging("[⚠️] Guardian returned no parseable response, defaulting to approved")
        return {"status": "approved", "allowed": True, "reason": "No guardian response parsed"}

    except Exception as e:
        abi_logging(f"[❌] Guardian validation failed: {e}")
        return {
            "status": "error",
            "allowed": False,
            "reason": f"Guardian error: {str(e)}",
        }


# ── Level 1: Gate decision (depends on both level 0 nodes) ──────

@agent.step(
    name="gate_decision",
    depends_on=["classify_query", "guardian_validate"],
    input_map={
        "triage": "$classify_query",
        "guardian": "$guardian_validate",
        "query": "$input.query",
    },
)
def gate_decision(triage, guardian, query):
    """Merge triage + guardian results and decide how to proceed."""
    guardian_status = guardian.get("status", "error")

    # Guardian failed (timeout, unreachable, etc.)
    if guardian_status == "error":
        abi_logging(f"[❌] Guardian failed: {guardian.get('reason', 'unknown')}")
        return {
            "action": "system_error",
            "message": "El sistema experimentó una falla en la validación de seguridad. Por favor reintente más tarde.",
            "guardian_reason": guardian.get("reason", ""),
        }

    # Guardian blocked (injection, reverse engineering, policy violation)
    if guardian_status == "blocked":
        abi_logging(f"[🛡️] Query blocked by guardian: {guardian.get('reason', '')}")
        return {
            "action": "blocked",
            "message": "Tu solicitud no puede ser procesada por razones de seguridad.",
            "guardian_reason": guardian.get("reason", ""),
        }

    # Guardian approved — proceed based on triage
    classification = triage.get("classification", "complex")

    # Clarification answer — deterministically route to planner with enriched query
    if classification == "clarification_response":
        original = triage.get("original_query", "")
        answer = triage.get("user_answer", query)
        enriched = (
            f"{original}\n\nUser clarification: {answer}"
            if original else answer
        )
        abi_logging("[🔁] Gate: clarification answer, re-routing to planner")
        return {"action": "call_planner", "query": enriched}

    if classification == "simple":
        abi_logging(f"[✅] Gate: simple query, responding directly")
        return {"action": "respond_direct", "query": query}
    else:
        abi_logging(f"[✅] Gate: complex query, calling planner")
        return {"action": "call_planner", "query": query}


# ── Level 2+: Planning pipeline (only runs if gate says call_planner) ──

@agent.step(
    name="call_planner",
    depends_on=["gate_decision"],
    input_map={
        "gate": "$gate_decision",
        "query": "$input.query",
        "context_id": "$input.context_id",
        "task_id": "$input.task_id",
    },
)
async def call_planner(gate, query, context_id, task_id):
    """Call Planner agent and return raw A2A results.

    Skips if gate_decision action is not 'call_planner'.
    """
    if gate.get("action") != "call_planner":
        return {"gate_passthrough": gate}

    # The gate may carry an enriched query (e.g. a clarification answer combined
    # with the original request). Prefer it over the raw input query.
    query = gate.get("query") or query

    abi_logging(f"[📞] Calling Planner: {query}")

    planner_card = await tool_find_agent.ainvoke({"query": "planner"})
    if not planner_card:
        raise ValueError("Could not find Planner agent")

    workflow = AgentInteractionFlow()
    node = InteractionFlowNode(
        task=query,
        source_agent_card=AGENT_CARD,
        target_agent_card=planner_card,
        node_key="planner",
        node_label="Planning Phase",
    )
    workflow.add_node(node)
    workflow.set_node_attributes(
        node.id, {"context_id": context_id, "task_id": task_id, "query": query}
    )
    workflow.set_source_card(AGENT_CARD)

    results = []
    async for chunk in workflow.run_workflow():
        results.append(chunk)
    return results


@agent.step(
    name="extract_plan",
    depends_on=["call_planner"],
    input_map={"planner_results": "$call_planner"},
)
def extract_plan(planner_results):
    """Extract execution plan from Planner results using A2AResponse."""
    # Gate passthrough — not a planner response
    if isinstance(planner_results, dict) and "gate_passthrough" in planner_results:
        return planner_results

    abi_logging(f"[🔍] extract_plan received {len(planner_results)} results")
    for i, r in enumerate(planner_results):
        parsed = A2AResponse.parse(r)
        abi_logging(f"  [{i}] {parsed}")

    needs_clarification, msg = A2AResponse.find_clarification(planner_results)
    if needs_clarification:
        return {"clarification": msg}

    plan = A2AResponse.find_plan(planner_results)
    if not plan:
        return {"error": "Could not generate execution plan"}

    abi_logging(f"[📋] Plan received with {len(plan.get('tasks', []))} tasks")
    return {"plan": plan}


@agent.step(
    name="build_workflow",
    depends_on=["extract_plan"],
    input_map={
        "plan_result": "$extract_plan",
        "context_id": "$input.context_id",
        "task_id": "$input.task_id",
    },
)
async def build_workflow(plan_result, context_id, task_id):
    """Build AgentInteractionFlow from the extracted plan.

    Handles three task types:
    - "execute": agent exists → add directly to workflow
    - "build_and_execute": no agent, tools exist → call builder first
    - "create_tools_and_execute": no agent, no tools → call builder first
    """
    # Gate passthrough or errors — pass through
    if "gate_passthrough" in plan_result or "clarification" in plan_result or "error" in plan_result:
        return plan_result

    plan = plan_result["plan"]
    workflow = AgentInteractionFlow()
    nodes = {}
    tasks = plan.get("tasks", [])
    ephemeral_agents = []

    abi_logging(f"[🔨] Creating workflow with {len(tasks)} tasks")

    # Find builder once (reused for all build tasks)
    builder_card = None
    needs_builder = any(
        t.get("type") in ("build_and_execute", "create_tools_and_execute")
        for t in tasks
    )
    if needs_builder:
        builder_card = await tool_find_agent.ainvoke({"query": "builder"})
        if not builder_card:
            return {"error": "Builder agent not found — cannot create ephemeral agents"}
        abi_logging(f"[🔧] Builder agent found: {builder_card.name}")

    for task in tasks:
        tid = task.get("task_id")
        desc = task.get("description", "")
        task_type = task.get("type", "execute")
        target = task.get("target", {})

        # Collect artifact keys from dependency targets
        dep_artifact_keys = []
        for dep_id in task.get("dependencies", []):
            dep_task = next((t for t in tasks if t.get("task_id") == dep_id), None)
            if dep_task and dep_task.get("target", {}).get("type") == "file":
                dep_tag = dep_task["target"]["tag"]
                # Key format matches what synthesize_and_report uploads
                dep_artifact_keys.append(dep_tag)

        if task_type == "execute":
            agents = task.get("agents", [])
            if not agents or not agents[0]:
                abi_logging(f"[⚠️] Task {tid}: no agent assigned, skipping")
                continue

            agent_dict = agents[0]
            target = (
                build_agent_card(agent_dict)[0]
                if isinstance(agent_dict, dict)
                else agent_dict
            )

            # Health check before adding to workflow
            target_url = get_agent_url(target)
            if target_url:
                health = await AbiAgent.check_health(target_url, target.name)
                if health.get("status") not in ("healthy",):
                    abi_logging(
                        f"[❌] Task {tid}: agent '{target.name}' unreachable "
                        f"({health.get('status')}), attempting cleanup"
                    )
                    # Deregister only if it's NOT an infrastructure agent
                    agent_name_lower = target.name.lower()
                    is_infra = any(infra in agent_name_lower for infra in INFRA_AGENTS)
                    if not is_infra:
                        try:
                            toolkit = MCPToolkit()
                            dereg_result = await toolkit.call("unregister_agent", agent_name=target.name)
                            if isinstance(dereg_result, dict) and dereg_result.get("success"):
                                abi_logging(f"[🗑️] Deregistered stale agent '{target.name}'")
                            else:
                                abi_logging(f"[⚠️] Deregister failed for '{target.name}': {dereg_result}")
                        except Exception as e:
                            abi_logging(f"[⚠️] Deregister error for '{target.name}': {e}")
                    else:
                        abi_logging(f"[🛡️] Skipping deregister for infrastructure agent '{target.name}'")
                    continue

            abi_logging(f"[✅] Task {tid}: execute → {target.name}")

        elif task_type in ("build_and_execute", "create_tools_and_execute"):
            builder_spec = task.get("builder_spec", {})
            # Pass artifact keys from dependencies and target tag
            if dep_artifact_keys:
                builder_spec["artifact_keys"] = dep_artifact_keys
            if target and target.get("tag"):
                builder_spec["target_tag"] = target["tag"]
            abi_logging(f"[🏗️] Task {tid}: {task_type} → calling builder (artifacts={dep_artifact_keys})")

            build_query = json.dumps({
                "task_id": tid,
                "task_type": task_type,
                "builder_spec": builder_spec,
                "description": desc,
            })

            build_flow = AgentInteractionFlow()
            build_node = InteractionFlowNode(
                task=build_query,
                source_agent_card=AGENT_CARD,
                target_agent_card=builder_card,
                node_key=f"build_{tid}",
                node_label=f"Build agent for {tid}",
            )
            build_flow.add_node(build_node)
            build_flow.set_node_attributes(
                build_node.id,
                {"context_id": context_id, "task_id": task_id, "query": build_query},
            )
            build_flow.set_source_card(AGENT_CARD)

            build_results = []
            async for chunk in build_flow.run_workflow():
                build_results.append(chunk)

            builder_response = A2AResponse.find_plan(build_results)
            if not builder_response:
                for resp in A2AResponse.from_results(build_results):
                    if resp.data:
                        builder_response = resp.data
                        break
                    if resp.text:
                        try:
                            parsed = json.loads(resp.text)
                            if isinstance(parsed, dict) and ("agent" in parsed or "agent_card" in parsed):
                                builder_response = parsed
                                break
                        except (json.JSONDecodeError, TypeError):
                            pass

            if not builder_response or builder_response.get("status") == "error":
                error_msg = (
                    builder_response.get("message", "Builder failed")
                    if builder_response
                    else "No response from builder"
                )
                abi_logging(f"[❌] Task {tid}: builder failed — {error_msg}")
                continue

            ephemeral_agent = builder_response.get("agent", {})
            ephemeral_card_data = builder_response.get("agent_card", {})

            if not ephemeral_card_data:
                abi_logging(f"[❌] Task {tid}: builder returned no agent card")
                continue

            target = build_agent_card(ephemeral_card_data)[0]
            ephemeral_agents.append(ephemeral_agent)

            abi_logging(
                f"[✅] Task {tid}: ephemeral agent '{ephemeral_agent.get('name')}' "
                f"ready at {ephemeral_agent.get('url')}"
            )
        else:
            abi_logging(f"[⚠️] Task {tid}: unknown type '{task_type}', skipping")
            continue

        abi_logging(f"[✅] Task {tid}: assigned to agent '{target.name}' at {get_agent_url(target)} with prompt {desc}")
        node = InteractionFlowNode(
            task=desc,
            source_agent_card=AGENT_CARD,
            target_agent_card=target,
            node_key=tid,
            node_label=f"{tid}: {desc[:40]}",
        )
        workflow.add_node(node)
        nodes[tid] = node
        workflow.set_node_attributes(
            node.id,
            {"task_id": task_id, "context_id": context_id, "query": desc},
        )

    for task in tasks:
        tid = task.get("task_id")
        for dep in task.get("dependencies", []):
            if dep in nodes and tid in nodes:
                workflow.add_edge(nodes[dep].id, nodes[tid].id)
                abi_logging(f"[🔗] Edge: {dep} → {tid}")

    workflow.set_source_card(AGENT_CARD)

    if workflow.is_empty():
        abi_logging("[⚠️] Workflow is empty — no agents could be assigned")
        return {"error": "No agents could be assigned to execute the plan. Please try a different request."}

    return {
        "workflow": workflow,
        "plan": plan,
        "ephemeral_agents": ephemeral_agents,
    }
