import asyncio
import json
from collections.abc import AsyncIterable

from abi_core.common import prompts
from abi_core.common.utils import abi_logging
from abi_core.common.workflow import Status
from abi_core.common.semantic_tools import tool_find_agent
from abi_core.agent.agent import AbiAgent
from abi_core.agent.agent_response import AgentResponse

from config import config


class AbiOrchestratorAgent(AbiAgent):
    """Orchestrator Agent — coordinates multi-agent workflows.

    The planning pipeline (call_planner -> extract_plan -> build_workflow)
    is registered as @agent.task() decorators in main.py and injected
    as self.tool_graph by AbiCore.

    build_workflow now handles three task types:
    - "execute": direct A2A call to existing agent
    - "build_and_execute": builder creates ephemeral agent, then execute
    - "create_tools_and_execute": builder creates tools + agent, then execute

    After workflow completion, ephemeral agents are destroyed.
    """

    def __init__(self):
        super().__init__(
            agent_name=config.AGENT_NAME,
            description=config.AGENT_DESCRIPTION,
            llm_config=config.LLM_CONFIG,
            tools=[tool_find_agent],
            system_prompt=prompts.ORCHESTRATOR_TOT_INSTRUCTIONS,
            content_types=['text', 'text/plain'],
        )

    async def _cleanup_ephemeral(self, ephemeral_agents):
        """Destroy ephemeral agent containers after task completion."""
        for agent_info in ephemeral_agents:
            name = agent_info.get("name", "")
            if not name or not agent_info.get("ephemeral"):
                continue

            abi_logging(f"[🗑️] Destroying ephemeral agent '{name}'")
            try:
                proc = await asyncio.create_subprocess_exec(
                    "docker", "rm", "-f", name,
                    stdout=asyncio.subprocess.PIPE,
                    stderr=asyncio.subprocess.PIPE,
                )
                await proc.communicate()
                if proc.returncode == 0:
                    abi_logging(f"[✅] Ephemeral agent '{name}' destroyed")
                else:
                    abi_logging(f"[⚠️] Could not destroy '{name}' (may already be gone)")
            except Exception as e:
                abi_logging(f"[⚠️] Cleanup error for '{name}': {e}")

    async def stream(
        self, query: str, context_id: str, task_id: str
    ) -> AsyncIterable[dict[str, any]]:
        """Orchestrate workflow execution using the task DAG."""

        abi_logging(f'[*] Orchestrator stream - context: {context_id}, task: {task_id}')
        abi_logging(f'[📝] Query: {query}')

        if not query:
            raise ValueError('Please provide a Query')

        ephemeral_agents = []

        try:
            if self.tool_graph is None:
                yield AgentResponse.error("No tool_graph configured")
                return

            # ── Phase 1: Planning pipeline (with heartbeat) ─────
            yield AgentResponse.status(
                "Planning...",
                agent=self.agent_name,
                context_id=context_id,
                task_id=task_id,
            )

            dag_coro = self.tool_graph.execute({
                "query": query,
                "context_id": context_id,
                "task_id": task_id,
            })
            dag_result, heartbeats = await self._run_with_heartbeat(
                dag_coro, context_id, task_id, "Planning in progress..."
            )
            for hb in heartbeats:
                yield hb

            if dag_result.get("failed_node"):
                error = dag_result.get("error", "Pipeline failed")
                abi_logging(f"[❌] DAG failed at {dag_result['failed_node']}: {error}")
                yield AgentResponse.error(error)
                return

            outputs = dag_result.get("node_outputs", {})
            build_result = outputs.get("build_workflow", {})

            if "clarification" in build_result:
                abi_logging("[❓] Forwarding clarification request to user")
                formatted = (
                    "🤔 **Necesito mas informacion para crear el mejor plan:**"
                    f"\n\n{build_result['clarification']}"
                )
                yield AgentResponse.input_required(formatted)
                return

            if "error" in build_result:
                yield AgentResponse.error(build_result["error"])
                return

            # ── Phase 2: Execute agent workflow ──────────────────
            workflow = build_result["workflow"]
            plan = build_result["plan"]
            ephemeral_agents = build_result.get("ephemeral_agents", [])

            abi_logging(
                f"[📋] Executing plan with {len(plan.get('tasks', []))} tasks"
                f" ({len(ephemeral_agents)} ephemeral)"
            )

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

            # ── Phase 3: Synthesize results (with heartbeat) ────
            if workflow.state == Status.COMPLETED:
                abi_logging(f"[✅] Workflow completed with {len(results)} results")

                synthesis_query = (
                    f"Synthesize the following workflow results:\n"
                    f"Plan: {json.dumps(plan, indent=2)}\n"
                    f"Results count: {len(results)}"
                )

                result_holder = {}

                async def _synthesize():
                    inputs = {"messages": [{"role": "user", "content": synthesis_query}]}
                    thread_config = {"configurable": {"thread_id": context_id}}
                    final = None
                    async for chunk in self.agent.astream(inputs, config=thread_config, stream_mode="updates"):
                        for _node, node_data in chunk.items():
                            if "messages" in node_data:
                                for msg in node_data["messages"]:
                                    if hasattr(msg, 'content') and msg.content:
                                        final = msg.content
                    result_holder['synthesis'] = final

                _, heartbeats = await self._run_with_heartbeat(
                    _synthesize(), context_id, task_id, "Synthesizing results..."
                )
                for hb in heartbeats:
                    yield hb

                yield AgentResponse.success(
                    result_holder.get('synthesis') or "Workflow completed successfully"
                )

        except Exception as e:
            abi_logging(f"[❌] Error in orchestration: {e}")
            yield AgentResponse.error(str(e))

        finally:
            # ── Phase 4: Cleanup ephemeral agents ───────────────
            if ephemeral_agents:
                abi_logging(f"[🧹] Cleaning up {len(ephemeral_agents)} ephemeral agent(s)")
                await self._cleanup_ephemeral(ephemeral_agents)
