"""
Agent Supervisor for Breakpilot

Agent lifecycle management, health monitoring, restart policies, load balancing.
"""

from typing import Dict, Optional, List, Any
from datetime import datetime, timezone, timedelta
import asyncio
import logging

from agent_core.sessions.heartbeat import HeartbeatMonitor
from agent_core.orchestrator.message_bus import (
    MessageBus,
    AgentMessage,
    MessagePriority,
)
from agent_core.orchestrator.supervisor_models import (
    AgentStatus,
    RestartPolicy,
    AgentInfo,
    AgentFactory,
)

logger = logging.getLogger(__name__)


class AgentSupervisor:
    """Supervises agents: lifecycle, health monitoring, restart policies, load balancing."""

    def __init__(
        self, message_bus: MessageBus,
        heartbeat_monitor: Optional[HeartbeatMonitor] = None,
        check_interval_seconds: int = 10
    ):
        self.bus = message_bus
        self.heartbeat = heartbeat_monitor or HeartbeatMonitor()
        self.check_interval = check_interval_seconds

        self.agents: Dict[str, AgentInfo] = {}
        self._factories: Dict[str, AgentFactory] = {}
        self._running = False
        self._health_task: Optional[asyncio.Task] = None

        # Set up heartbeat timeout handler
        self.heartbeat.on_timeout = self._handle_agent_timeout

    async def start(self) -> None:
        """Starts the supervisor"""
        self._running = True
        await self.heartbeat.start_monitoring()

        # Start health check loop
        self._health_task = asyncio.create_task(self._health_check_loop())

        logger.info("Agent supervisor started")

    async def stop(self) -> None:
        """Stops the supervisor and all agents"""
        self._running = False

        # Stop health check
        if self._health_task:
            self._health_task.cancel()
            try:
                await self._health_task
            except asyncio.CancelledError:
                pass

        # Stop heartbeat monitor
        await self.heartbeat.stop_monitoring()

        # Stop all agents
        for agent_id in list(self.agents.keys()):
            await self.stop_agent(agent_id)

        logger.info("Agent supervisor stopped")

    def register_factory(
        self,
        agent_type: str,
        factory: AgentFactory
    ) -> None:
        """
        Registers a factory function for creating agents.

        Args:
            agent_type: Type of agent this factory creates
            factory: Async function that creates agent instances
        """
        self._factories[agent_type] = factory
        logger.debug(f"Registered factory for agent type: {agent_type}")

    async def register_agent(
        self,
        agent_id: str,
        agent_type: str,
        restart_policy: RestartPolicy = RestartPolicy.ON_FAILURE,
        max_restarts: int = 3,
        capacity: int = 10,
        metadata: Optional[Dict[str, Any]] = None
    ) -> AgentInfo:
        """
        Registers a new agent with the supervisor.

        Args:
            agent_id: Unique agent identifier
            agent_type: Type of agent
            restart_policy: When to restart the agent
            max_restarts: Maximum restart attempts
            capacity: Max concurrent tasks
            metadata: Additional agent metadata

        Returns:
            AgentInfo for the registered agent
        """
        if agent_id in self.agents:
            logger.warning(f"Agent {agent_id} already registered")
            return self.agents[agent_id]

        info = AgentInfo(
            agent_id=agent_id,
            agent_type=agent_type,
            restart_policy=restart_policy,
            max_restarts=max_restarts,
            capacity=capacity,
            metadata=metadata or {}
        )

        self.agents[agent_id] = info
        self.heartbeat.register(agent_id, agent_type)

        logger.info(f"Registered agent: {agent_id} ({agent_type})")

        return info

    async def start_agent(self, agent_id: str) -> bool:
        """
        Starts a registered agent.

        Args:
            agent_id: The agent to start

        Returns:
            True if started successfully
        """
        if agent_id not in self.agents:
            logger.error(f"Cannot start unregistered agent: {agent_id}")
            return False

        info = self.agents[agent_id]

        if info.status == AgentStatus.RUNNING:
            logger.warning(f"Agent {agent_id} is already running")
            return True

        info.status = AgentStatus.STARTING

        try:
            # If we have a factory, create the agent
            if info.agent_type in self._factories:
                factory = self._factories[info.agent_type]
                await factory(agent_id)

            info.status = AgentStatus.RUNNING
            info.started_at = datetime.now(timezone.utc)
            info.last_activity = info.started_at

            # Subscribe to message bus
            await self.bus.subscribe(
                agent_id,
                self._create_message_handler(agent_id)
            )

            logger.info(f"Started agent: {agent_id}")
            return True

        except Exception as e:
            info.status = AgentStatus.ERROR
            info.error_count += 1
            logger.error(f"Failed to start agent {agent_id}: {e}")
            return False

    async def stop_agent(self, agent_id: str) -> bool:
        """
        Stops a running agent.

        Args:
            agent_id: The agent to stop

        Returns:
            True if stopped successfully
        """
        if agent_id not in self.agents:
            return False

        info = self.agents[agent_id]
        info.status = AgentStatus.STOPPING

        try:
            # Unsubscribe from message bus
            await self.bus.unsubscribe(agent_id)

            # Unregister from heartbeat
            self.heartbeat.unregister(agent_id)

            info.status = AgentStatus.STOPPED
            logger.info(f"Stopped agent: {agent_id}")
            return True

        except Exception as e:
            info.status = AgentStatus.ERROR
            logger.error(f"Error stopping agent {agent_id}: {e}")
            return False

    async def restart_agent(self, agent_id: str) -> bool:
        """
        Restarts an agent.

        Args:
            agent_id: The agent to restart

        Returns:
            True if restarted successfully
        """
        if agent_id not in self.agents:
            return False

        info = self.agents[agent_id]

        # Check restart count
        if info.restart_count >= info.max_restarts:
            logger.error(
                f"Agent {agent_id} exceeded max restarts "
                f"({info.restart_count}/{info.max_restarts})"
            )
            await self._escalate_agent_failure(agent_id)
            return False

        info.status = AgentStatus.RESTARTING
        info.restart_count += 1

        logger.info(
            f"Restarting agent {agent_id} "
            f"(attempt {info.restart_count}/{info.max_restarts})"
        )

        # Stop and start
        await self.stop_agent(agent_id)
        await asyncio.sleep(1)  # Brief pause
        return await self.start_agent(agent_id)

    async def _handle_agent_timeout(
        self,
        session_id: str,
        agent_type: str
    ) -> None:
        """Handles agent heartbeat timeout"""
        # Find the agent by session/ID
        agent_id = session_id  # In our case, session_id == agent_id

        if agent_id not in self.agents:
            return

        info = self.agents[agent_id]
        info.status = AgentStatus.ERROR
        info.error_count += 1

        logger.warning(f"Agent {agent_id} timed out (heartbeat)")

        # Apply restart policy
        if info.restart_policy == RestartPolicy.NEVER:
            await self._escalate_agent_failure(agent_id)
        elif info.restart_policy == RestartPolicy.ON_FAILURE:
            if info.restart_count < info.max_restarts:
                await self.restart_agent(agent_id)
            else:
                await self._escalate_agent_failure(agent_id)
        elif info.restart_policy == RestartPolicy.ALWAYS:
            await self.restart_agent(agent_id)

    async def _escalate_agent_failure(self, agent_id: str) -> None:
        """Escalates an agent failure to the alert system"""
        info = self.agents.get(agent_id)
        if not info:
            return

        await self.bus.publish(AgentMessage(
            sender="supervisor",
            receiver="alert-agent",
            message_type="agent_failure",
            payload={
                "agent_id": agent_id,
                "agent_type": info.agent_type,
                "error_count": info.error_count,
                "restart_count": info.restart_count,
                "last_activity": info.last_activity.isoformat() if info.last_activity else None
            },
            priority=MessagePriority.CRITICAL
        ))

        logger.error(f"Escalated agent failure: {agent_id}")

    def _create_message_handler(self, agent_id: str):
        """Creates a message handler that updates agent activity"""
        async def handler(message: AgentMessage):
            if agent_id in self.agents:
                self.agents[agent_id].last_activity = datetime.now(timezone.utc)
                # Heartbeat on activity
                self.heartbeat.beat(agent_id)
            return None
        return handler

    async def _health_check_loop(self) -> None:
        """Periodic health check loop"""
        while self._running:
            try:
                await asyncio.sleep(self.check_interval)
                await self._run_health_checks()
            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error(f"Error in health check loop: {e}")

    async def _run_health_checks(self) -> None:
        """Runs health checks on all agents"""
        now = datetime.now(timezone.utc)

        for agent_id, info in list(self.agents.items()):
            if info.status != AgentStatus.RUNNING:
                continue

            # Check for stale agents (no activity for 5 minutes)
            if info.last_activity:
                idle_time = now - info.last_activity
                if idle_time > timedelta(minutes=5):
                    logger.warning(
                        f"Agent {agent_id} has been idle for "
                        f"{idle_time.total_seconds():.0f}s"
                    )

    # Load Balancing

    def get_available_agent(
        self,
        agent_type: str,
        strategy: str = "least_loaded"
    ) -> Optional[str]:
        """
        Gets an available agent of the specified type.

        Args:
            agent_type: Type of agent needed
            strategy: Load balancing strategy

        Returns:
            Agent ID or None if none available
        """
        candidates = [
            info for info in self.agents.values()
            if info.agent_type == agent_type and info.is_available()
        ]

        if not candidates:
            return None

        if strategy == "least_loaded":
            # Pick agent with lowest load
            best = min(candidates, key=lambda a: a.utilization())
        elif strategy == "round_robin":
            # Simple round-robin (just pick first available)
            best = candidates[0]
        else:
            best = candidates[0]

        return best.agent_id

    def acquire_capacity(self, agent_id: str) -> bool:
        """
        Acquires capacity from an agent.

        Args:
            agent_id: The agent to acquire from

        Returns:
            True if capacity was acquired
        """
        if agent_id not in self.agents:
            return False

        info = self.agents[agent_id]
        if not info.is_available():
            return False

        info.current_load += 1
        return True

    def release_capacity(self, agent_id: str) -> None:
        """
        Releases capacity back to an agent.

        Args:
            agent_id: The agent to release to
        """
        if agent_id in self.agents:
            info = self.agents[agent_id]
            info.current_load = max(0, info.current_load - 1)

    # Status and Metrics

    def get_agent_status(self, agent_id: str) -> Optional[Dict[str, Any]]:
        """Gets status information for an agent"""
        if agent_id not in self.agents:
            return None

        info = self.agents[agent_id]
        return {
            "agent_id": info.agent_id,
            "agent_type": info.agent_type,
            "status": info.status.value,
            "current_task": info.current_task,
            "started_at": info.started_at.isoformat() if info.started_at else None,
            "last_activity": info.last_activity.isoformat() if info.last_activity else None,
            "error_count": info.error_count,
            "restart_count": info.restart_count,
            "utilization": info.utilization(),
            "is_healthy": info.is_healthy(),
            "is_available": info.is_available()
        }

    def get_all_status(self) -> Dict[str, Dict[str, Any]]:
        """Gets status for all agents"""
        return {
            agent_id: self.get_agent_status(agent_id)
            for agent_id in self.agents
        }

    def get_metrics(self) -> Dict[str, Any]:
        """Gets aggregate metrics"""
        total = len(self.agents)
        running = sum(
            1 for a in self.agents.values()
            if a.status == AgentStatus.RUNNING
        )
        healthy = sum(1 for a in self.agents.values() if a.is_healthy())
        available = sum(1 for a in self.agents.values() if a.is_available())

        total_capacity = sum(a.capacity for a in self.agents.values())
        total_load = sum(a.current_load for a in self.agents.values())

        return {
            "total_agents": total,
            "running_agents": running,
            "healthy_agents": healthy,
            "available_agents": available,
            "total_capacity": total_capacity,
            "total_load": total_load,
            "overall_utilization": total_load / max(total_capacity, 1),
            "by_type": self._metrics_by_type()
        }

    def _metrics_by_type(self) -> Dict[str, Dict[str, int]]:
        """Gets metrics grouped by agent type"""
        by_type: Dict[str, Dict[str, int]] = {}

        for info in self.agents.values():
            if info.agent_type not in by_type:
                by_type[info.agent_type] = {
                    "total": 0,
                    "running": 0,
                    "healthy": 0
                }

            by_type[info.agent_type]["total"] += 1
            if info.status == AgentStatus.RUNNING:
                by_type[info.agent_type]["running"] += 1
            if info.is_healthy():
                by_type[info.agent_type]["healthy"] += 1

        return by_type
