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'],
        )

    def _record_error(self, context_id: str, error_type: str, message: str):
        """Record an error in session context for next request awareness."""
        ctx = self.get_session_context(context_id)
        ctx[f"last_error_{error_type}"] = message
        if not hasattr(self, '_conversation_history'):
            self._conversation_history = {}
        self._conversation_history[context_id] = ctx
        abi_logging(f"[📝] Error recorded in session {context_id}: {error_type}")

    async def _cleanup_ephemeral(self, ephemeral_agents):
        """Destroy ephemeral agent containers and deregister from semantic layer."""
        from abi_core.common.container_runtime import destroy_container
        from abi_core.common.semantic_tools import mcp_toolkit

        for agent_info in ephemeral_agents:
            name = agent_info.get("name", "")
            if not name or not agent_info.get("ephemeral"):
                continue

            # Deregister from semantic layer
            try:
                await mcp_toolkit.call("unregister_agent", agent_name=name)
                abi_logging(f"[📋] Deregistered ephemeral agent '{name}' from semantic layer")
            except Exception as e:
                abi_logging(f"[⚠️] Could not deregister '{name}': {e}")

            # Destroy container
            await destroy_container(name)

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

        The DAG now has a parallel first phase:
        - classify_query (triage) + guardian_validate (security) run in parallel
        - gate_decision merges both and decides: respond_direct, call_planner, blocked, system_error
        """

        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: Triage + Guardian + Planning pipeline (DAG) ──
            yield AgentResponse.status(
                "Analyzing request...",
                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, "Processing..."
            )
            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}")
                self._record_error(context_id, "dag_failed", error)
                yield AgentResponse.error(error)
                return

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

            # ── Check gate decision first ────────────────────────
            gate = outputs.get("gate_decision", {})
            action = gate.get("action", "") if isinstance(gate, dict) else ""
            abi_logging(f"[🚦] Gate decision: action={action}, outputs_keys={list(outputs.keys())}")

            if action == "system_error":
                self._record_error(context_id, "guardian_failed", gate.get("message", ""))
                yield AgentResponse.error(gate.get("message", "System error"))
                return

            if action == "blocked":
                self._record_error(context_id, "blocked", gate.get("message", ""))
                yield AgentResponse.error(gate.get("message", "Request blocked"))
                return

            if action == "respond_direct":
                # Simple query — respond with orchestrator's own agent (has checkpointer memory)
                yield AgentResponse.status(
                    "Responding...",
                    agent=self.agent_name,
                    context_id=context_id,
                    task_id=task_id,
                )

                result_holder = {}

                async def _respond_direct():
                    inputs = {"messages": [{"role": "user", "content": 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['response'] = final

                _, heartbeats = await self._run_with_heartbeat(
                    _respond_direct(), context_id, task_id, "Thinking..."
                )
                for hb in heartbeats:
                    yield hb

                yield AgentResponse.success(result_holder.get('response') or "No response generated")
                return

            # ── action == "call_planner" — check planning results ──
            build_result = outputs.get("build_workflow", {})

            # Gate passthrough means planner was skipped (shouldn't happen here but safety check)
            if "gate_passthrough" in build_result:
                gate_pass = build_result["gate_passthrough"]
                yield AgentResponse.error(gate_pass.get("message", "Unexpected gate passthrough"))
                return

            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:
                self._record_error(context_id, "plan_error", build_result["error"])
                yield AgentResponse.error(build_result["error"])
                return

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

            if not workflow or workflow.is_empty():
                msg = "No agents could be assigned to execute the plan. Please try a different request."
                self._record_error(context_id, "empty_workflow", msg)
                yield AgentResponse.error(msg)
                return

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

                # Extract artifact URLs from agent responses
                artifacts = []
                for r in results:
                    try:
                        resp = A2AResponse.parse(r)
                        if resp and resp.data:
                            for art in resp.data.get("uploaded_artifacts", []):
                                artifacts.append(art)
                    except Exception:
                        pass

                # Generate pre-signed download URLs
                if artifacts:
                    try:
                        from abi_core.common.artifact_store import ArtifactStore
                        store = ArtifactStore()
                        for art in artifacts:
                            key = art.get("key") or f"{art.get('agent', 'unknown')}/{art.get('filename', '')}"
                            art["download_url"] = await store.get_url(key, expires=3600)
                        abi_logging(f"[📎] {len(artifacts)} artifact(s) with download URLs")
                    except Exception as e:
                        abi_logging(f"[⚠️] Could not generate download URLs: {e}")

                # Build synthesis prompt with artifacts info
                artifacts_section = ""
                if artifacts:
                    artifacts_section = "\nGenerated artifacts:\n"
                    for art in artifacts:
                        fname = art.get("filename", "unknown")
                        url = art.get("download_url", art.get("url", ""))
                        artifacts_section += f"  - {fname}: {url}\n"

                synthesis_query = (
                    f"Synthesize the following workflow results:\n"
                    f"Plan: {json.dumps(plan, indent=2)}\n"
                    f"Results count: {len(results)}"
                    f"{artifacts_section}\n"
                    f"Include download links for any generated files in your response."
                )

                from abi_core.agent.llm_provider import invoke as llm_invoke

                async def _synthesize():
                    return await llm_invoke(
                        config.LLM_CONFIG, synthesis_query,
                        thread_id=context_id,
                    )

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

                # Append artifact links if LLM didn't include them
                final_response = synthesis or "Workflow completed successfully"
                if artifacts and "download" not in final_response.lower():
                    final_response += "\n\n📎 **Generated files:**\n"
                    for art in artifacts:
                        fname = art.get("filename", "unknown")
                        url = art.get("download_url", art.get("url", ""))
                        final_response += f"- [{fname}]({url})\n"

                yield AgentResponse.success(final_response)

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

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