from __future__ import annotations

"""Multi-agent orchestrator."""

import asyncio
import json
from dataclasses import dataclass, field
from typing import Any

from agentrspy.core import (
    AgentError,
    AgentOutput,
    AgentTrait,
    CompletionRequest,
    LlmProvider,
    Message,
    Usage,
)

from .communication import EventBus, OrchestratorEvent
from .graph import AgentGraph
from .shared_memory import SharedConversation


@dataclass(slots=True)
class RoutingStrategy:
    """Routing strategy for multi-agent runs."""

    kind: str
    names: list[str] | None = None
    llm: LlmProvider | None = None
    agents: list[str] | None = None
    max_turns: int | None = None
    graph: AgentGraph | None = None

    @classmethod
    def sequential(cls, names: list[str]) -> "RoutingStrategy":
        return cls(kind="sequential", names=names)

    @classmethod
    def parallel(cls, names: list[str]) -> "RoutingStrategy":
        return cls(kind="parallel", names=names)

    @classmethod
    def supervisor(
        cls, llm: LlmProvider, agents: list[str], max_turns: int
    ) -> "RoutingStrategy":
        return cls(kind="supervisor", llm=llm, agents=agents, max_turns=max_turns)

    @classmethod
    def graph_strategy(cls, graph: AgentGraph) -> "RoutingStrategy":
        return cls(kind="graph", graph=graph)


RoutingStrategy.Sequential = RoutingStrategy.sequential  # type: ignore[attr-defined]
RoutingStrategy.Parallel = RoutingStrategy.parallel  # type: ignore[attr-defined]
RoutingStrategy.Supervisor = RoutingStrategy.supervisor  # type: ignore[attr-defined]
RoutingStrategy.Graph = RoutingStrategy.graph_strategy  # type: ignore[attr-defined]


@dataclass(slots=True)
class _SharedAgent:
    agent: AgentTrait
    lock: asyncio.Lock = field(default_factory=asyncio.Lock)


@dataclass(slots=True)
class MultiAgentOrchestrator:
    """Multi-agent orchestrator."""

    agents: dict[str, _SharedAgent]
    routing: RoutingStrategy
    shared_memory: SharedConversation | None = None
    event_bus: EventBus | None = None

    @classmethod
    def builder(cls) -> "MultiAgentOrchestratorBuilder":
        return MultiAgentOrchestratorBuilder()

    async def run(self, input_text: str) -> AgentOutput:
        if self.routing.kind == "sequential":
            return await self._run_sequential(
                input_text, list(self.routing.names or [])
            )
        if self.routing.kind == "parallel":
            return await self._run_parallel(input_text, list(self.routing.names or []))
        if self.routing.kind == "supervisor":
            assert (
                self.routing.llm is not None
                and self.routing.agents is not None
                and self.routing.max_turns is not None
            )
            return await self._run_supervisor(
                input_text,
                self.routing.llm,
                self.routing.agents,
                self.routing.max_turns,
            )
        if self.routing.kind == "graph":
            assert self.routing.graph is not None
            return await self._run_graph(input_text, self.routing.graph)
        raise AgentError.invalid_configuration(
            f"unsupported routing strategy: {self.routing.kind}"
        )

    def add_agent_boxed(self, name: str, agent: AgentTrait) -> None:
        self.agents[name] = _SharedAgent(agent=agent)

    async def _run_sequential(self, input_text: str, names: list[str]) -> AgentOutput:
        current_input = input_text
        final_output: AgentOutput | None = None
        for name in names:
            output = await self._run_named_agent(name, current_input)
            current_input = output.text
            final_output = output
        if final_output is None:
            raise AgentError.invalid_configuration("no agents configured")
        return final_output

    async def _run_parallel(self, input_text: str, names: list[str]) -> AgentOutput:
        async def run_one(name: str) -> tuple[str, AgentOutput]:
            return name, await self._run_named_agent(name, input_text)

        results = await asyncio.gather(*(run_one(name) for name in names))
        messages = []
        text_parts = []
        for name, output in results:
            text_parts.append(f"[{name}]\n{output.text}")
            messages.extend(output.messages)
        return AgentOutput(
            text="\n\n".join(text_parts),
            steps=1,
            usage=Usage(),
            messages=messages,
            metadata={},
        )

    async def _run_supervisor(
        self, input_text: str, llm: LlmProvider, agents: list[str], max_turns: int
    ) -> AgentOutput:
        agent_lines = "\n".join(f"- {name}" for name in agents)
        context = (
            "You are a supervisor. Available agents:\n"
            f"{agent_lines}\n\n"
            f'Return JSON {{"agent": "name"}} for the best agent to handle the task. Task: {input_text}'
        )
        for _ in range(max_turns):
            response = await llm.complete(
                CompletionRequest(
                    messages=[Message.user(context)],
                    tools=None,
                    model="",
                    temperature=0.0,
                    max_tokens=256,
                    stream=False,
                    system=None,
                    extra={},
                )
            )
            try:
                choice = json.loads(response.message.text_content())
            except json.JSONDecodeError:
                choice = None
            if isinstance(choice, dict) and isinstance(choice.get("agent"), str):
                return await self._run_named_agent(choice["agent"], input_text)
            context = f"{context}\nPrevious response was invalid JSON: {response.message.text_content()}"
        raise AgentError.max_steps_reached(max_turns)

    async def _run_graph(self, input_text: str, graph: AgentGraph) -> AgentOutput:
        current = graph.entry()
        current_input = input_text
        while True:
            output = await self._run_named_agent(current, current_input)
            next_name = graph.next(current, output)
            if next_name is None:
                return output
            current = next_name
            current_input = output.text

    async def _run_named_agent(self, name: str, input_text: str) -> AgentOutput:
        shared_agent = self.agents.get(name)
        if shared_agent is None:
            raise AgentError.agent_not_found(name)
        async with shared_agent.lock:
            output = await shared_agent.agent.run(input_text)
        if self.shared_memory is not None and output.messages:
            await self.shared_memory.add(name, output.messages[-1])
        if self.event_bus is not None:
            await self.event_bus.publish(OrchestratorEvent(agent=name, output=output))
        return output


@dataclass(slots=True)
class MultiAgentOrchestratorBuilder:
    """Builder for `MultiAgentOrchestrator`."""

    agents: dict[str, _SharedAgent] = field(default_factory=dict)
    order: list[str] = field(default_factory=list)
    routing_strategy: RoutingStrategy | None = None
    shared_memory_value: SharedConversation | None = None
    event_bus_value: EventBus | None = None

    def add_agent(
        self, name: str, agent: AgentTrait
    ) -> "MultiAgentOrchestratorBuilder":
        self.order.append(name)
        self.agents[name] = _SharedAgent(agent=agent)
        return self

    def add_agent_boxed(
        self, name: str, agent: AgentTrait
    ) -> "MultiAgentOrchestratorBuilder":
        return self.add_agent(name, agent)

    def routing(self, routing: RoutingStrategy) -> "MultiAgentOrchestratorBuilder":
        self.routing_strategy = routing
        return self

    def shared_memory(
        self, shared_memory: SharedConversation
    ) -> "MultiAgentOrchestratorBuilder":
        self.shared_memory_value = shared_memory
        return self

    def event_bus(self, event_bus: EventBus) -> "MultiAgentOrchestratorBuilder":
        self.event_bus_value = event_bus
        return self

    def build(self) -> MultiAgentOrchestrator:
        routing = self.routing_strategy or RoutingStrategy.sequential(list(self.order))
        return MultiAgentOrchestrator(
            agents=self.agents,
            routing=routing,
            shared_memory=self.shared_memory_value,
            event_bus=self.event_bus_value,
        )
