"""
Agent Orchestrator

Manages the lifecycle of all agents in the system.
Handles spawning, monitoring, and coordination.
"""

import asyncio
import contextlib
from collections.abc import Callable
from datetime import UTC, datetime
from typing import Any
from uuid import UUID

import structlog

from roboco.agents.base import Agent
from roboco.models import AgentRole, AgentStatus, Team

logger = structlog.get_logger()


class Orchestrator:
    """
    Central orchestrator for all RoboCo agents.

    Responsibilities:
    - Spawn and stop agents
    - Monitor agent health
    - Route messages between agents
    - Handle agent failures and restarts
    """

    def __init__(self) -> None:
        """Initialize the orchestrator."""
        self._agents: dict[UUID, Agent] = {}
        self._agent_tasks: dict[UUID, asyncio.Task] = {}
        self._running = False
        self._monitor_task: asyncio.Task | None = None

        self.log = logger.bind(component="orchestrator")

    @property
    def agents(self) -> dict[UUID, Agent]:
        """Get all registered agents."""
        return self._agents.copy()

    @property
    def active_agents(self) -> list[Agent]:
        """Get all active agents."""
        return [
            a for a in self._agents.values() if a.state.status == AgentStatus.ACTIVE
        ]

    @property
    def idle_agents(self) -> list[Agent]:
        """Get all idle agents."""
        return [a for a in self._agents.values() if a.state.status == AgentStatus.IDLE]

    # =========================================================================
    # LIFECYCLE
    # =========================================================================

    async def start(self) -> None:
        """Start the orchestrator."""
        if self._running:
            self.log.warning("Orchestrator already running")
            return

        self.log.info("Starting orchestrator")
        self._running = True

        # Start health monitor
        self._monitor_task = asyncio.create_task(self._health_monitor())

    async def stop(self) -> None:
        """Stop the orchestrator and all agents."""
        if not self._running:
            return

        self.log.info("Stopping orchestrator")
        self._running = False

        # Stop health monitor
        if self._monitor_task:
            self._monitor_task.cancel()
            with contextlib.suppress(asyncio.CancelledError):
                await self._monitor_task

        # Stop all agents
        await self.stop_all_agents()

        self.log.info("Orchestrator stopped")

    # =========================================================================
    # AGENT MANAGEMENT
    # =========================================================================

    def register_agent(self, agent: Agent) -> None:
        """
        Register an agent with the orchestrator.

        Args:
            agent: The agent to register
        """
        if agent.id in self._agents:
            self.log.warning("Agent already registered", agent_id=str(agent.id))
            return

        self._agents[agent.id] = agent
        self.log.info(
            "Agent registered",
            agent_id=str(agent.id),
            agent_name=agent.name,
            agent_role=agent.role.value,
        )

    def unregister_agent(self, agent_id: UUID) -> None:
        """
        Unregister an agent.

        Args:
            agent_id: ID of the agent to unregister
        """
        if agent_id not in self._agents:
            return

        agent = self._agents.pop(agent_id)
        self.log.info(
            "Agent unregistered", agent_id=str(agent_id), agent_name=agent.name
        )

    async def spawn_agent(self, agent: Agent) -> None:
        """
        Spawn an agent (register and start).

        Args:
            agent: The agent to spawn
        """
        self.register_agent(agent)
        await agent.start()
        self.log.info("Agent spawned", agent_id=str(agent.id), agent_name=agent.name)

    async def stop_agent(self, agent_id: UUID) -> None:
        """
        Stop a specific agent.

        Args:
            agent_id: ID of the agent to stop
        """
        if agent_id not in self._agents:
            self.log.warning("Agent not found", agent_id=str(agent_id))
            return

        agent = self._agents[agent_id]
        await agent.stop()
        self.log.info("Agent stopped", agent_id=str(agent_id), agent_name=agent.name)

    async def stop_all_agents(self) -> None:
        """Stop all registered agents."""
        self.log.info("Stopping all agents", count=len(self._agents))

        # Stop all agents concurrently
        await asyncio.gather(
            *[agent.stop() for agent in self._agents.values()],
            return_exceptions=True,
        )

    async def restart_agent(self, agent_id: UUID) -> None:
        """
        Restart an agent.

        Args:
            agent_id: ID of the agent to restart
        """
        if agent_id not in self._agents:
            self.log.warning("Agent not found", agent_id=str(agent_id))
            return

        agent = self._agents[agent_id]
        self.log.info("Restarting agent", agent_id=str(agent_id), agent_name=agent.name)

        await agent.stop()
        await asyncio.sleep(1)  # Brief pause
        await agent.start()

    # =========================================================================
    # QUERYING
    # =========================================================================

    def get_agent(self, agent_id: UUID) -> Agent | None:
        """Get an agent by ID."""
        return self._agents.get(agent_id)

    def get_agent_by_slug(self, slug: str) -> Agent | None:
        """Get an agent by slug."""
        for agent in self._agents.values():
            if agent.config.slug == slug:
                return agent
        return None

    def get_agents_by_role(self, role: AgentRole) -> list[Agent]:
        """Get all agents with a specific role."""
        return [a for a in self._agents.values() if a.role == role]

    def get_agents_by_team(self, team: Team) -> list[Agent]:
        """Get all agents in a specific team."""
        return [a for a in self._agents.values() if a.team == team]

    def get_cell_agents(self, team: Team) -> dict[str, list[Agent]]:
        """
        Get agents organized by role for a team/cell.

        Returns:
            Dict with keys: developers, qa, pm, documenter
        """
        team_agents = self.get_agents_by_team(team)
        return {
            "developers": [a for a in team_agents if a.role == AgentRole.DEVELOPER],
            "qa": [a for a in team_agents if a.role == AgentRole.QA],
            "pm": [a for a in team_agents if a.role == AgentRole.CELL_PM],
            "documenter": [a for a in team_agents if a.role == AgentRole.DOCUMENTER],
        }

    # =========================================================================
    # HEALTH MONITORING
    # =========================================================================

    async def _health_monitor(self) -> None:
        """
        Monitor agent health periodically.

        Checks for:
        - Unresponsive agents
        - Agents with errors
        - Agents that need restart
        """
        while self._running:
            try:
                await self._check_agent_health()
                await asyncio.sleep(30)  # Check every 30 seconds
            except asyncio.CancelledError:
                break
            except Exception as e:
                self.log.error("Error in health monitor", error=str(e))
                await asyncio.sleep(5)

    async def _check_agent_health(self) -> None:
        """Check health of all agents."""
        now = datetime.now(UTC)

        for agent in self._agents.values():
            # Check for errors
            if agent.state.error:
                self.log.warning(
                    "Agent has error",
                    agent_id=str(agent.id),
                    error=agent.state.error,
                )

            # Check for inactivity (5 minutes)
            minutes_in_seconds = 300
            if agent.state.last_activity:
                inactive_seconds = (now - agent.state.last_activity).total_seconds()
                if inactive_seconds > minutes_in_seconds and agent.is_running:
                    self.log.warning(
                        "Agent inactive",
                        agent_id=str(agent.id),
                        inactive_seconds=inactive_seconds,
                    )

    def get_health_status(self) -> dict[str, Any]:
        """
        Get overall health status of all agents.

        Returns:
            Health status summary
        """
        total = len(self._agents)
        by_status: dict[str, int] = {}
        errors: list[dict[str, str]] = []

        for agent in self._agents.values():
            status = agent.state.status.value
            by_status[status] = by_status.get(status, 0) + 1

            if agent.state.error:
                errors.append(
                    {
                        "agent_id": str(agent.id),
                        "agent_name": agent.name,
                        "error": agent.state.error,
                    }
                )

        return {
            "total_agents": total,
            "by_status": by_status,
            "errors": errors,
            "healthy": len(errors) == 0,
        }

    # =========================================================================
    # CELL MANAGEMENT
    # =========================================================================

    async def spawn_cell(
        self,
        team: Team,
        agent_factory: Callable[[Team], list[Agent]],
    ) -> list[Agent]:
        """
        Spawn all agents for a cell.

        Args:
            team: The team/cell to spawn
            agent_factory: Factory function that creates agents for the team

        Returns:
            List of spawned agents
        """
        self.log.info("Spawning cell", team=team.value)

        agents = agent_factory(team)
        for agent in agents:
            await self.spawn_agent(agent)

        return agents

    async def stop_cell(self, team: Team) -> None:
        """
        Stop all agents in a cell.

        Args:
            team: The team/cell to stop
        """
        self.log.info("Stopping cell", team=team.value)

        team_agents = self.get_agents_by_team(team)
        await asyncio.gather(
            *[self.stop_agent(a.id) for a in team_agents],
            return_exceptions=True,
        )

    # =========================================================================
    # SERIALIZATION
    # =========================================================================

    def to_dict(self) -> dict[str, Any]:
        """Get orchestrator status as dictionary."""
        return {
            "running": self._running,
            "total_agents": len(self._agents),
            "agents": [a.to_dict() for a in self._agents.values()],
            "health": self.get_health_status(),
        }


# =============================================================================
# SINGLETON HOLDER
# =============================================================================


class _OrchestratorHolder:
    """Holder class for singleton orchestrator instance."""

    instance: Orchestrator | None = None


def get_orchestrator() -> Orchestrator:
    """Get or create the global orchestrator instance."""
    if _OrchestratorHolder.instance is None:
        _OrchestratorHolder.instance = Orchestrator()
    return _OrchestratorHolder.instance


async def start_orchestrator() -> Orchestrator:
    """Start the global orchestrator."""
    orchestrator = get_orchestrator()
    await orchestrator.start()
    return orchestrator


async def stop_orchestrator() -> None:
    """Stop the global orchestrator."""
    if _OrchestratorHolder.instance:
        await _OrchestratorHolder.instance.stop()
        _OrchestratorHolder.instance = None
