"""Graceful shutdown handling for Oclawma.

This module provides signal handlers for SIGTERM/SIGINT to ensure
in-progress tasks complete before exit, preventing data corruption.
"""

from __future__ import annotations

import asyncio
import json
import logging
import signal
import time
from collections.abc import Awaitable, Callable
from dataclasses import dataclass, field
from enum import Enum
from pathlib import Path
from typing import Any

logger = logging.getLogger(__name__)


class ShutdownState(Enum):
    """Current shutdown state."""

    RUNNING = "running"
    SHUTDOWN_REQUESTED = "shutdown_requested"
    SHUTTING_DOWN = "shutting_down"
    SHUTDOWN_COMPLETE = "shutdown_complete"


@dataclass
class JobInfo:
    """Information about an active job."""

    job_id: str
    job_type: str
    start_time: float
    metadata: dict[str, Any] = field(default_factory=dict)


@dataclass
class QueueState:
    """State of the job queue for persistence."""

    version: int = 1
    timestamp: float = field(default_factory=time.time)
    pending_jobs: list[dict[str, Any]] = field(default_factory=list)
    active_jobs: list[dict[str, Any]] = field(default_factory=list)
    completed_jobs: list[dict[str, Any]] = field(default_factory=list)

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary for serialization."""
        return {
            "version": self.version,
            "timestamp": self.timestamp,
            "pending_jobs": self.pending_jobs,
            "active_jobs": self.active_jobs,
            "completed_jobs": self.completed_jobs,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> QueueState:
        """Create from dictionary."""
        return cls(
            version=data.get("version", 1),
            timestamp=data.get("timestamp", time.time()),
            pending_jobs=data.get("pending_jobs", []),
            active_jobs=data.get("active_jobs", []),
            completed_jobs=data.get("completed_jobs", []),
        )


class ShutdownManager:
    """Manager for graceful shutdown handling.

    Handles SIGTERM/SIGINT signals and ensures:
    - Active jobs complete before exit (with timeout)
    - Queue state is saved before exiting
    - Resources are cleaned up properly
    """

    def __init__(
        self,
        timeout_seconds: float = 30.0,
        state_file: Path | str | None = None,
        loop: asyncio.AbstractEventLoop | None = None,
    ):
        """Initialize the shutdown manager.

        Args:
            timeout_seconds: Maximum time to wait for jobs to complete
            state_file: Path to save queue state on shutdown
            loop: Event loop to use (defaults to asyncio.get_event_loop())
        """
        self.timeout_seconds = timeout_seconds
        self.state_file = Path(state_file) if state_file else None
        self._loop = loop
        self._state = ShutdownState.RUNNING
        self._active_jobs: dict[str, JobInfo] = {}
        self._pending_jobs: list[dict[str, Any]] = []
        self._completed_jobs: list[dict[str, Any]] = []
        self._cleanup_handlers: list[Callable[[], None]] = []
        self._async_cleanup_handlers: list[Callable[[], Awaitable[Any]]] = []
        self._shutdown_event = asyncio.Event()
        self._signal_handlers_installed = False

    @property
    def state(self) -> ShutdownState:
        """Get current shutdown state."""
        return self._state

    @property
    def is_shutting_down(self) -> bool:
        """Check if shutdown has been requested."""
        return self._state in (ShutdownState.SHUTDOWN_REQUESTED, ShutdownState.SHUTTING_DOWN)

    @property
    def active_job_count(self) -> int:
        """Get number of active jobs."""
        return len(self._active_jobs)

    @property
    def pending_job_count(self) -> int:
        """Get number of pending jobs."""
        return len(self._pending_jobs)

    def install_signal_handlers(self) -> None:
        """Install signal handlers for SIGTERM and SIGINT."""
        if self._signal_handlers_installed:
            return

        try:
            # Only works on Unix-like systems
            loop = self._get_loop()
            for sig in (signal.SIGTERM, signal.SIGINT):
                loop.add_signal_handler(sig, self._signal_handler)
            self._signal_handlers_installed = True
            logger.debug("Signal handlers installed for SIGTERM and SIGINT")
        except (NotImplementedError, ValueError) as e:
            # Windows or already handled - set up fallback
            logger.warning(f"Could not install signal handlers: {e}")
            # Fallback: register signal handlers via signal module
            try:
                signal.signal(signal.SIGTERM, self._signal_handler_fallback)
                signal.signal(signal.SIGINT, self._signal_handler_fallback)
                self._signal_handlers_installed = True
            except (ValueError, OSError) as e2:
                logger.warning(f"Fallback signal handler installation failed: {e2}")

    def uninstall_signal_handlers(self) -> None:
        """Uninstall signal handlers."""
        if not self._signal_handlers_installed:
            return

        try:
            loop = self._get_loop()
            for sig in (signal.SIGTERM, signal.SIGINT):
                loop.remove_signal_handler(sig)
        except (NotImplementedError, ValueError):
            pass

        try:
            signal.signal(signal.SIGTERM, signal.SIG_DFL)
            signal.signal(signal.SIGINT, signal.SIG_DFL)
        except (ValueError, OSError):
            pass

        self._signal_handlers_installed = False

    def _get_loop(self) -> asyncio.AbstractEventLoop:
        """Get the event loop."""
        if self._loop is None:
            try:
                self._loop = asyncio.get_running_loop()
            except RuntimeError:
                self._loop = asyncio.new_event_loop()
                asyncio.set_event_loop(self._loop)
        return self._loop

    def _signal_handler(self) -> None:
        """Handle SIGTERM/SIGINT signals (asyncio version)."""
        logger.info("Shutdown signal received")
        self._state = ShutdownState.SHUTDOWN_REQUESTED
        self._shutdown_event.set()

        # Schedule graceful shutdown
        loop = self._get_loop()
        loop.create_task(self._graceful_shutdown())

    def _signal_handler_fallback(self, signum: int, frame: Any) -> None:
        """Fallback signal handler using signal module."""
        sig_name = "SIGTERM" if signum == signal.SIGTERM else "SIGINT"
        logger.info(f"Shutdown signal received ({sig_name})")
        self._state = ShutdownState.SHUTDOWN_REQUESTED
        self._shutdown_event.set()

        # Can't schedule async from signal handler, so just trigger sync shutdown
        # The main loop should check is_shutting_down periodically

    async def _graceful_shutdown(self) -> None:
        """Perform graceful shutdown."""
        if self._state == ShutdownState.SHUTTING_DOWN:
            return

        self._state = ShutdownState.SHUTTING_DOWN
        logger.info(f"Starting graceful shutdown (timeout: {self.timeout_seconds}s)")
        logger.info(f"Active jobs: {self.active_job_count}, Pending jobs: {self.pending_job_count}")

        # Save queue state first
        await self._save_queue_state()

        # Wait for active jobs to complete
        if self._active_jobs:
            await self._wait_for_jobs()

        # Run cleanup handlers
        await self._run_cleanup_handlers()

        self._state = ShutdownState.SHUTDOWN_COMPLETE
        logger.info("Graceful shutdown complete")

    async def _wait_for_jobs(self) -> None:
        """Wait for active jobs to complete with timeout."""
        start_time = time.time()

        while self._active_jobs and (time.time() - start_time) < self.timeout_seconds:
            logger.info(f"Waiting for {self.active_job_count} active job(s)...")
            try:
                await asyncio.wait_for(
                    self._wait_for_all_jobs(),
                    timeout=min(5.0, self.timeout_seconds - (time.time() - start_time)),
                )
            except asyncio.TimeoutError:
                continue

        if self._active_jobs:
            logger.warning(
                f"Shutdown timeout reached. {self.active_job_count} job(s) still active."
            )
            # Cancel remaining jobs
            for job_id in list(self._active_jobs.keys()):
                logger.warning(f"Job {job_id} will be terminated")

    async def _wait_for_all_jobs(self) -> None:
        """Wait for all active jobs to complete."""
        while self._active_jobs:
            await asyncio.sleep(0.1)

    async def _save_queue_state(self) -> None:
        """Save the current queue state to file."""
        if not self.state_file:
            return

        try:
            state = QueueState(
                timestamp=time.time(),
                pending_jobs=self._pending_jobs.copy(),
                active_jobs=[
                    {
                        "job_id": job.job_id,
                        "job_type": job.job_type,
                        "start_time": job.start_time,
                        "metadata": job.metadata,
                    }
                    for job in self._active_jobs.values()
                ],
                completed_jobs=self._completed_jobs.copy(),
            )

            # Ensure parent directory exists
            self.state_file.parent.mkdir(parents=True, exist_ok=True)

            # Write atomically
            temp_file = self.state_file.with_suffix(".tmp")
            temp_file.write_text(json.dumps(state.to_dict(), indent=2))
            temp_file.rename(self.state_file)

            logger.info(f"Queue state saved to {self.state_file}")
        except Exception as e:
            logger.error(f"Failed to save queue state: {e}")

    async def _run_cleanup_handlers(self) -> None:
        """Run all registered cleanup handlers."""
        # Run sync handlers
        for handler in self._cleanup_handlers:
            try:
                handler()
            except Exception as e:
                logger.error(f"Cleanup handler failed: {e}")

        # Run async handlers
        for handler in self._async_cleanup_handlers:
            try:
                await handler()  # type: ignore[misc]
            except Exception as e:
                logger.error(f"Async cleanup handler failed: {e}")

    def register_job(
        self,
        job_id: str,
        job_type: str,
        metadata: dict[str, Any] | None = None,
    ) -> None:
        """Register an active job.

        Args:
            job_id: Unique identifier for the job
            job_type: Type of job (e.g., 'subagent', 'task')
            metadata: Optional metadata about the job
        """
        if self.is_shutting_down:
            raise RuntimeError("Cannot register new jobs during shutdown")

        self._active_jobs[job_id] = JobInfo(
            job_id=job_id,
            job_type=job_type,
            start_time=time.time(),
            metadata=metadata or {},
        )
        logger.debug(f"Job registered: {job_id} ({job_type})")

    def unregister_job(self, job_id: str, success: bool = True) -> None:
        """Unregister a completed job.

        Args:
            job_id: Unique identifier for the job
            success: Whether the job completed successfully
        """
        if job_id in self._active_jobs:
            job_info = self._active_jobs.pop(job_id)
            self._completed_jobs.append(
                {
                    "job_id": job_id,
                    "job_type": job_info.job_type,
                    "start_time": job_info.start_time,
                    "end_time": time.time(),
                    "success": success,
                    "metadata": job_info.metadata,
                }
            )
            logger.debug(f"Job unregistered: {job_id} (success={success})")

    def add_pending_job(self, job_data: dict[str, Any]) -> None:
        """Add a job to the pending queue.

        Args:
            job_data: Job data to store
        """
        self._pending_jobs.append(job_data)
        logger.debug(f"Pending job added: {job_data.get('job_id', 'unknown')}")

    def remove_pending_job(self, job_id: str) -> dict[str, Any] | None:
        """Remove and return a pending job by ID.

        Args:
            job_id: ID of the job to remove

        Returns:
            The job data if found, None otherwise
        """
        for i, job in enumerate(self._pending_jobs):
            if job.get("job_id") == job_id:
                return self._pending_jobs.pop(i)
        return None

    def register_cleanup_handler(self, handler: Callable[[], None]) -> None:
        """Register a synchronous cleanup handler.

        Args:
            handler: Function to call during cleanup
        """
        self._cleanup_handlers.append(handler)

    def register_async_cleanup_handler(
        self,
        handler: Callable[[], Awaitable[Any]],
    ) -> None:
        """Register an asynchronous cleanup handler.

        Args:
            handler: Async function to call during cleanup
        """
        self._async_cleanup_handlers.append(handler)

    async def wait_for_shutdown(self) -> None:
        """Wait until shutdown is requested.

        This can be used in the main event loop to block until shutdown.
        """
        await self._shutdown_event.wait()

    def load_queue_state(self) -> QueueState | None:
        """Load queue state from file.

        Returns:
            The loaded queue state, or None if no state file exists
        """
        if not self.state_file or not self.state_file.exists():
            return None

        try:
            data = json.loads(self.state_file.read_text())
            state = QueueState.from_dict(data)
            logger.info(f"Queue state loaded from {self.state_file}")
            return state
        except Exception as e:
            logger.error(f"Failed to load queue state: {e}")
            return None

    def clear_queue_state(self) -> None:
        """Clear the saved queue state file."""
        if self.state_file and self.state_file.exists():
            try:
                self.state_file.unlink()
                logger.info(f"Queue state cleared: {self.state_file}")
            except Exception as e:
                logger.error(f"Failed to clear queue state: {e}")

    def get_status(self) -> dict[str, Any]:
        """Get current shutdown manager status.

        Returns:
            Dictionary with current status information
        """
        return {
            "state": self._state.value,
            "is_shutting_down": self.is_shutting_down,
            "active_jobs": self.active_job_count,
            "pending_jobs": self.pending_job_count,
            "completed_jobs": len(self._completed_jobs),
            "timeout_seconds": self.timeout_seconds,
            "state_file": str(self.state_file) if self.state_file else None,
            "signal_handlers_installed": self._signal_handlers_installed,
        }


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


def get_shutdown_manager() -> ShutdownManager:
    """Get the global shutdown manager instance.

    Returns:
        The global ShutdownManager instance
    """
    global _shutdown_manager
    if _shutdown_manager is None:
        _shutdown_manager = ShutdownManager()
    return _shutdown_manager


def set_shutdown_manager(manager: ShutdownManager) -> None:
    """Set the global shutdown manager instance.

    Args:
        manager: The ShutdownManager instance to use
    """
    global _shutdown_manager
    _shutdown_manager = manager


def reset_shutdown_manager() -> None:
    """Reset the global shutdown manager instance."""
    global _shutdown_manager
    _shutdown_manager = None
