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

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 a2a.types import AgentCard
from config import AGENT_CARD

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


@agent.task(
    name="call_planner",
    input_map={
        "query": "$input.query",
        "context_id": "$input.context_id",
        "task_id": "$input.task_id",
    },
)
async def call_planner(query, context_id, task_id):
    """Call Planner agent and return raw A2A results."""
    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."""
    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
    """
    if "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 = []  # Track for cleanup

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

        if task_type == "execute":
            # ── Direct execution: agent already exists ──────────
            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 creates ephemeral agent first ───────────
            builder_spec = task.get("builder_spec", {})

            abi_logging(
                f"[🏗️] Task {tid}: {task_type} → calling builder"
            )

            # Call builder via A2A with the builder_spec as payload
            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)

            # Execute builder and collect results
            build_results = []
            async for chunk in build_flow.run_workflow():
                build_results.append(chunk)

            # Parse builder response
            builder_response = A2AResponse.find_plan(build_results)
            if not builder_response:
                # Try to get text/data from 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

            # Extract ephemeral agent card from builder response
            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

        # Add node to workflow
        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},
        )

    # Create edges from dependencies
    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)
    return {
        "workflow": workflow,
        "plan": plan,
        "ephemeral_agents": ephemeral_agents,
    }


agent.run(AbiOrchestratorAgent())
