"""
RouterOrchestrator — dynamic routing mode.

Design notes:
  - A Router Agent (or function) examines each incoming task and decides
    which specialist agent should handle it
  - Supports both LLM-based routing (the router is itself an agent) and
    rule-based routing (a simple function)
  - Optionally combines with the DiscoveryService for capability-based routing
"""

from __future__ import annotations

import logging
from collections.abc import Awaitable, Callable

from nexus.core.agent import Agent
from nexus.core.task import Task, TaskResult
from nexus.protocol.discovery import DiscoveryService

logger = logging.getLogger("nexus.orchestration.router")

# A routing function takes a task and returns the agent_id to handle it
RouteFunc = Callable[[Task], Awaitable[str | None]]


class RouterOrchestrator:
    """
    Dynamic routing orchestration — a router decides per-task which agent runs.

    Three routing strategies:
    1. Rule-based: a user-defined function maps tasks to agents
    2. Capability-based: uses DiscoveryService to match task requirements
    3. LLM-based: a router agent decides using natural language reasoning

    Usage:
        router = RouterOrchestrator(discovery=discovery)

        # Rule-based routing
        router.add_route("coding", agent_id=coder.id)
        router.add_route("writing", agent_id=writer.id)
        result = await router.route_and_execute(task)

        # Capability-based routing (automatic)
        task.required_capabilities = AgentCapability(coding=0.9)
        result = await router.route_and_execute(task)

        # LLM-based routing
        router = RouterOrchestrator(
            router_agent=triage_agent,
            agents={"coder": coder, "writer": writer},
        )
        result = await router.route_and_execute(task)
    """

    def __init__(
        self,
        discovery: DiscoveryService | None = None,
        router_agent: Agent | None = None,
        agents: dict[str, Agent] | None = None,
        default_agent_id: str | None = None,
    ) -> None:
        self.discovery = discovery or DiscoveryService()
        self.router_agent = router_agent
        self._agents = agents or {}
        self._rules: dict[str, str] = {}  # task_type → agent_id
        self._custom_route_func: RouteFunc | None = None
        self.default_agent_id = default_agent_id

    def add_agent(self, name: str, agent: Agent) -> None:
        """Register an agent for routing."""
        self._agents[name] = agent
        self.discovery.register(agent)

    def add_route(self, task_type: str, agent_id: str) -> None:
        """Add a rule-based route: tasks of this type go to this agent."""
        self._rules[task_type] = agent_id

    def set_route_func(self, func: RouteFunc) -> None:
        """Set a custom routing function."""
        self._custom_route_func = func

    async def resolve_agent(self, task: Task) -> Agent | None:
        """
        Determine which agent should handle a task.

        Priority:
        1. Custom route function
        2. Rule-based routing (by task_type)
        3. LLM-based routing (via router_agent)
        4. Capability-based routing (via discovery)
        5. Default agent
        """
        available = list(self._agents.keys())
        logger.info(
            f"🔶 Router resolving task '{task.name}' | "
            f"available agents: {available} | "
            f"task_type='{task.task_type}'"
        )

        # 1. Custom function
        if self._custom_route_func:
            agent_id = await self._custom_route_func(task)
            if agent_id:
                agent = self._find_agent(agent_id)
                if agent:
                    logger.debug(f"Custom route: task '{task.name}' → agent '{agent.name}'")
                    return agent

        # 2. Rule-based
        if task.task_type in self._rules:
            agent_id = self._rules[task.task_type]
            agent = self._find_agent(agent_id)
            if agent:
                logger.debug(f"Rule route: task type '{task.task_type}' → agent '{agent.name}'")
                return agent

        # 3. LLM-based
        if self.router_agent:
            logger.info(f"🤖 Using LLM router to decide agent for: '{task.description[:80]}...'")
            agent = await self._llm_route(task)
            if agent:
                logger.info(f"✅ LLM router decision: '{task.name}' → agent '{agent.name}'")
                return agent

        # 4. Capability-based
        if task.required_capabilities:
            best_id = self.discovery.find_best_agent(task.required_capabilities)
            if best_id:
                agent = self._find_agent(best_id)
                if agent:
                    logger.debug(f"Capability route: task '{task.name}' → agent '{agent.name}'")
                    return agent

        # 5. Default
        if self.default_agent_id:
            return self._find_agent(self.default_agent_id)

        logger.warning(f"No agent found for task '{task.name}'")
        return None

    async def route_and_execute(self, task: Task) -> TaskResult:
        """Resolve the best agent and execute the task."""
        agent = await self.resolve_agent(task)
        if not agent:
            task.fail("No suitable agent found by router")
            raise RuntimeError(f"Router could not find an agent for task '{task.name}'")

        logger.info(
            f"🚀 Executing task '{task.name}' with agent '{agent.name}' "
            f"(id={agent.id})"
        )
        result = await agent.run(task)
        logger.info(
            f"✅ Task '{task.name}' completed by '{agent.name}' "
            f"({len(result.content)} chars)"
        )
        return result

    async def _llm_route(self, task: Task) -> Agent | None:
        """
        Use the router agent (LLM) to decide which specialist handles the task.

        The router agent is given a description of available agents and
        asked to pick the best one.
        """
        if not self.router_agent:
            return None

        # Build a prompt listing available agents
        agent_descriptions = []
        for name, agent in self._agents.items():
            agent_descriptions.append(
                f"- **{name}** (id={agent.id}): {agent.instructions[:100]}..."
            )

        routing_task = Task(
            name="route_decision",
            description=(
                f"You are a routing agent. Given the following task, decide which "
                f"specialist agent should handle it.\n\n"
                f"**Task:** {task.description}\n\n"
                f"**Available agents:**\n" +
                "\n".join(agent_descriptions) +
                f"\n\nRespond with ONLY the agent name (one of: {list(self._agents.keys())}). "
                f"No explanation needed."
            ),
        )

        try:
            result = await self.router_agent.run(routing_task)
            chosen_name = result.content.strip().lower()
            logger.info(f"🤖 LLM router raw response: '{result.content.strip()}'")

            # Try to match the response to an agent name
            for name, agent in self._agents.items():
                if name.lower() in chosen_name or chosen_name in name.lower():
                    logger.info(f"🎯 LLM router matched: '{name}' for task '{task.name}'")
                    return agent

            logger.warning(
                f"⚠️ LLM router response '{chosen_name}' didn't match any agent "
                f"(available: {list(self._agents.keys())})"
            )
        except Exception as exc:
            logger.error(f"LLM routing failed: {exc}")

        return None

    def _find_agent(self, agent_id: str) -> Agent | None:
        """Look up an agent by ID or name."""
        # Check by ID in agents dict values
        for agent in self._agents.values():
            if agent.id == agent_id:
                return agent
        # Check by name
        if agent_id in self._agents:
            return self._agents[agent_id]
        # Check discovery service
        return self.discovery.get_agent(agent_id)
