#!/usr/bin/env python3
"""Orchestrator Agent — Generated by ABI-Core

DAG with parallel triage + guardian gate:

  Level 0 (parallel):  classify_query | guardian_validate
  Level 1 (merge):     gate_decision
  Level 2 (conditional): call_planner → extract_plan → build_workflow

classify_query and guardian_validate run in parallel (no deps between them).
gate_decision waits for both and decides: respond_direct, call_planner, blocked, or system_error.
"""

import json

from orchestrator import AbiOrchestratorAgent
from web_interface import OrchestratorWebinterface
from abi_core.agent import AbiCore
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
from abi_core.common.workflow import AgentInteractionFlow, InteractionFlowNode
from abi_core.common import prompts
from a2a.types import AgentCard
from config import AGENT_CARD, config

agent = AbiCore(
    web_interface_cls=OrchestratorWebinterface,
    interface_name="Orchestrator Web Interface",
)


# ── Level 0: Parallel triage + security ─────────────────────────

@agent.task(
    name="classify_query",
    input_map={"query": "$input.query"},
)
async def classify_query(query):
    """Classify query as simple or complex using the orchestrator's LLM."""
    from abi_core.agent.llm_provider import invoke
    from abi_core.common.utils import clean_llm_json

    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"}


@agent.task(
    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.task(
    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")

    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.task(
    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}

    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.task(
    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.task(
    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 = (
                AgentCard(**agent_dict)
                if isinstance(agent_dict, dict)
                else agent_dict
            )
            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 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 = AgentCard(**ephemeral_card_data)
            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 {target.url} 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,
    }


agent.run(AbiOrchestratorAgent())
