"""Routing strategies for the Orchestrator."""

from __future__ import annotations

from abc import ABC, abstractmethod
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from koan.agents.base import BaseAgent
    from koan.agents.types import RunContext
    from koan.llm.client import LLMClient
    from koan.llm.config import LLMConfig


class Router(ABC):
    """Abstract base for routing strategies."""

    @abstractmethod
    async def route(
        self,
        input_text: str,
        agents: dict[str, BaseAgent],
        *,
        context: RunContext | None = None,
    ) -> str | None:
        """Return the name of the agent to handle this input, or None."""
        ...


class RuleRouter(Router):
    """Route based on keyword matching rules.

    Usage::

        router = RuleRouter({
            "weather": "weather_agent",
            "stock": "finance_agent",
            "code": "code_agent",
        })
    """

    def __init__(self, rules: dict[str, str], *, default: str | None = None) -> None:
        self.rules = rules
        self.default = default

    async def route(
        self,
        input_text: str,
        agents: dict[str, BaseAgent],
        *,
        context: RunContext | None = None,
    ) -> str | None:
        lower = input_text.lower()
        for keyword, agent_name in self.rules.items():
            if keyword.lower() in lower and agent_name in agents:
                return agent_name
        return self.default if self.default and self.default in agents else None


class LLMRouter(Router):
    """Route by asking an LLM to classify the input.

    The LLM is prompted with the available agent names and descriptions,
    and asked to return the name of the best agent.

    Usage::

        router = LLMRouter(client=llm_client, config=llm_config)
    """

    def __init__(self, *, client: LLMClient, config: LLMConfig) -> None:
        self._client = client
        self._config = config

    async def route(
        self,
        input_text: str,
        agents: dict[str, BaseAgent],
        *,
        context: RunContext | None = None,
    ) -> str | None:
        from koan.types import Message

        agent_list = ", ".join(agents.keys())
        prompt = (
            f"You are a routing classifier. Given the user input, respond with ONLY "
            f"the name of the best agent to handle it. Available agents: {agent_list}\n\n"
            f"User input: {input_text}\n\nAgent name:"
        )

        messages = [Message(role="user", content=prompt)]
        response = await self._client.complete(self._config, messages)

        chosen = response.text.strip().strip("'\"")
        return chosen if chosen in agents else None
