"""Multi-agent orchestration – sequential and parallel execution modes."""

from __future__ import annotations

import asyncio
from typing import Any

from thryve.core.agent.agent import Agent
from thryve.core.agent.models import StopReason, TurnResult
from thryve.core.context.models import Message
from thryve.utils import get_logger

logger = get_logger("agent.multi")


class AgentLaborMarket:
    """Registry of available agents."""

    def __init__(self) -> None:
        self._agents: dict[str, Agent] = {}

    def register(self, agent: Agent) -> None:
        self._agents[agent.name] = agent

    def get(self, name: str) -> Agent | None:
        return self._agents.get(name)

    def list_agents(self) -> list[Agent]:
        return list(self._agents.values())


class MultiAgentOrchestrator:
    """Coordinate multiple agents working on a task.

    Modes
    -----
    * **sequential** – A main agent drives the workflow and may delegate
      sub-tasks to worker agents.
    * **parallel** – The task is split into independent sub-tasks that are
      executed concurrently by different agents.
    """

    def __init__(self) -> None:
        self.labor_market = AgentLaborMarket()
        self.main_agent: Agent | None = None

    # ------------------------------------------------------------------
    # Configuration
    # ------------------------------------------------------------------

    def register_agent(self, agent: Agent) -> None:
        """Register *agent* in the labor market."""
        self.labor_market.register(agent)

    def set_main_agent(self, agent: Agent) -> None:
        """Designate *agent* as the primary orchestrator."""
        self.main_agent = agent
        self.labor_market.register(agent)

    # ------------------------------------------------------------------
    # Execution
    # ------------------------------------------------------------------

    async def execute(
        self,
        task: str,
        mode: str = "sequential",
        **kwargs: Any,
    ) -> TurnResult:
        """Dispatch *task* using the chosen *mode*.

        Parameters
        ----------
        task:
            Natural-language description of the task.
        mode:
            ``"sequential"`` or ``"parallel"``.
        """
        if mode == "sequential":
            return await self._execute_sequential(task, **kwargs)
        elif mode == "parallel":
            return await self._execute_parallel(task, **kwargs)
        else:
            raise ValueError(f"Unknown mode: {mode!r}")

    # ------------------------------------------------------------------
    # Sequential
    # ------------------------------------------------------------------

    async def _execute_sequential(self, task: str, **kwargs: Any) -> TurnResult:
        """Main agent runs the task; it can optionally delegate to workers."""
        if self.main_agent is None:
            agents = self.labor_market.list_agents()
            if not agents:
                raise RuntimeError("No agents registered")
            self.main_agent = agents[0]

        return await self.main_agent.run(task, **kwargs)

    # ------------------------------------------------------------------
    # Parallel
    # ------------------------------------------------------------------

    async def _execute_parallel(self, task: str, **kwargs: Any) -> TurnResult:
        """Run all registered agents on *task* concurrently.

        Returns a merged :class:`TurnResult` that concatenates individual
        responses.
        """
        agents = self.labor_market.list_agents()
        if not agents:
            raise RuntimeError("No agents registered")

        results: list[TurnResult] = await asyncio.gather(
            *(agent.run(task, **kwargs) for agent in agents)
        )

        # Merge: concatenate final responses, collect all steps.
        combined_text = "\n\n---\n\n".join(
            f"[{agents[i].name}]\n{r.final_response}" for i, r in enumerate(results)
        )
        all_steps = [step for r in results for step in r.steps]

        return TurnResult(
            final_response=combined_text,
            steps=all_steps,
            stop_reason=StopReason.COMPLETED,
        )
