"""Graceful shutdown handling for the agent."""

from __future__ import annotations

import asyncio
import signal
from collections.abc import Callable, Coroutine
from contextlib import asynccontextmanager
from typing import TYPE_CHECKING, Any

import structlog

if TYPE_CHECKING:
    from collections.abc import AsyncGenerator

log = structlog.get_logger(__name__)

ShutdownCallback = Callable[[], Coroutine[Any, Any, None]]


class ShutdownManager:
    """
    Manages graceful shutdown of the agent.

    Handles SIGTERM and SIGINT signals, coordinates cleanup of async tasks,
    and ensures resources are properly released.
    """

    def __init__(self, shutdown_timeout: float = 30.0) -> None:
        """
        Initialize shutdown manager.

        Args:
            shutdown_timeout: Maximum time to wait for graceful shutdown
        """
        self._shutdown_timeout = shutdown_timeout
        self._shutdown_event = asyncio.Event()
        self._callbacks: list[tuple[str, ShutdownCallback]] = []
        self._shutting_down = False
        self._original_handlers: dict[signal.Signals, Any] = {}

    @property
    def is_shutting_down(self) -> bool:
        """Check if shutdown is in progress."""
        return self._shutting_down

    @property
    def shutdown_event(self) -> asyncio.Event:
        """Get the shutdown event for waiting."""
        return self._shutdown_event

    def register_callback(self, name: str, callback: ShutdownCallback) -> None:
        """
        Register a callback to be called during shutdown.

        Callbacks are called in reverse order of registration (LIFO).

        Args:
            name: Descriptive name for logging
            callback: Async function to call during shutdown
        """
        self._callbacks.append((name, callback))
        log.debug("registered_shutdown_callback", callback_name=name)

    def unregister_callback(self, name: str) -> None:
        """Remove a registered callback by name."""
        self._callbacks = [(n, cb) for n, cb in self._callbacks if n != name]

    async def wait_for_shutdown(self) -> None:
        """Wait until shutdown is signaled."""
        await self._shutdown_event.wait()

    def trigger_shutdown(self) -> None:
        """Programmatically trigger shutdown."""
        if not self._shutting_down:
            log.info("shutdown_triggered_programmatically")
            self._shutdown_event.set()

    def _signal_handler(self, sig: signal.Signals) -> None:
        """Handle shutdown signals."""
        if self._shutting_down:
            log.warning("shutdown_already_in_progress", signal=sig.name)
            return

        log.info("shutdown_signal_received", signal=sig.name)
        self._shutdown_event.set()

    def install_signal_handlers(self) -> None:
        """Install signal handlers for SIGTERM and SIGINT."""
        loop = asyncio.get_running_loop()

        for sig in (signal.SIGTERM, signal.SIGINT):
            try:
                self._original_handlers[sig] = signal.getsignal(sig)
                loop.add_signal_handler(
                    sig,
                    lambda s: self._signal_handler(s),
                    sig,
                )
                log.debug("installed_signal_handler", signal=sig.name)
            except NotImplementedError:
                # Windows doesn't support add_signal_handler
                signal.signal(
                    sig,
                    lambda s, f, sig=sig: self._signal_handler(sig),  # type: ignore[misc]
                )
                log.debug("installed_signal_handler_fallback", signal=sig.name)

    def remove_signal_handlers(self) -> None:
        """Remove installed signal handlers."""
        loop = asyncio.get_running_loop()

        for sig, original in self._original_handlers.items():
            try:
                loop.remove_signal_handler(sig)
            except NotImplementedError:
                signal.signal(sig, original or signal.SIG_DFL)

        self._original_handlers.clear()

    async def execute_shutdown(self) -> None:
        """Execute the shutdown sequence."""
        if self._shutting_down:
            return

        self._shutting_down = True
        log.info("shutdown_started", callback_count=len(self._callbacks))

        # Execute callbacks in reverse order (LIFO)
        for name, callback in reversed(self._callbacks):
            try:
                log.debug("executing_shutdown_callback", callback_name=name)
                await asyncio.wait_for(
                    callback(),
                    timeout=self._shutdown_timeout / len(self._callbacks)
                    if self._callbacks
                    else self._shutdown_timeout,
                )
                log.debug("shutdown_callback_completed", callback_name=name)
            except TimeoutError:
                log.error("shutdown_callback_timeout", callback_name=name)
            except Exception as e:
                log.error("shutdown_callback_error", callback_name=name, error=str(e))

        self.remove_signal_handlers()
        log.info("shutdown_completed")

    @asynccontextmanager
    async def managed_lifecycle(self) -> AsyncGenerator[ShutdownManager, None]:
        """
        Context manager for managing the agent lifecycle.

        Installs signal handlers on entry and executes shutdown on exit.
        """
        self.install_signal_handlers()
        try:
            yield self
        finally:
            await self.execute_shutdown()


# Global shutdown manager instance
_shutdown_manager: ShutdownManager | None = None


def get_shutdown_manager(shutdown_timeout: float = 30.0) -> ShutdownManager:
    """Get or create the global shutdown manager."""
    global _shutdown_manager
    if _shutdown_manager is None:
        _shutdown_manager = ShutdownManager(shutdown_timeout)
    return _shutdown_manager
