"""
Health and metrics HTTP server.

Provides endpoints for:
- /healthz - Liveness probe
- /ready - Readiness probe (matches Kubernetes convention)
- /metrics - Prometheus metrics
- /state - State transfer (GET/POST) for pod migration
- /prestop - PreStop hook handler for graceful shutdown
"""

import logging
from typing import TYPE_CHECKING, Callable, Awaitable

from aiohttp import web

from dory.health.probes import LivenessProbe, ReadinessProbe
from dory.utils.errors import DoryHealthError

if TYPE_CHECKING:
    from dory.metrics.collector import MetricsCollector

logger = logging.getLogger(__name__)

# Type aliases for callbacks
StateGetter = Callable[[], dict]
StateRestorer = Callable[[dict], Awaitable[None]]
PreStopHandler = Callable[[], Awaitable[None]]


class HealthServer:
    """
    HTTP server for health probes, metrics, and state transfer.

    Runs on a separate port from the main application.
    Provides endpoints required by Dory Orchestrator for:
    - Health probes (liveness/readiness)
    - Prometheus metrics
    - State transfer during pod migration
    - PreStop hook for graceful shutdown
    """

    def __init__(
        self,
        port: int = 8080,
        health_path: str = "/healthz",
        ready_path: str = "/ready",  # Changed from /readyz to match Orchestrator
        metrics_path: str = "/metrics",
        metrics_collector: "MetricsCollector | None" = None,
        state_getter: StateGetter | None = None,
        state_restorer: StateRestorer | None = None,
        prestop_handler: PreStopHandler | None = None,
    ):
        """
        Initialize health server.

        Args:
            port: Port to listen on
            health_path: Path for liveness probe
            ready_path: Path for readiness probe
            metrics_path: Path for Prometheus metrics
            metrics_collector: Optional metrics collector for /metrics endpoint
            state_getter: Callback to get processor state for /state GET
            state_restorer: Callback to restore processor state for /state POST
            prestop_handler: Callback for /prestop PreStop hook
        """
        self._port = port
        self._health_path = health_path
        self._ready_path = ready_path
        self._metrics_path = metrics_path
        self._metrics_collector = metrics_collector
        self._state_getter = state_getter
        self._state_restorer = state_restorer
        self._prestop_handler = prestop_handler

        self._liveness = LivenessProbe()
        self._readiness = ReadinessProbe()

        self._app: web.Application | None = None
        self._runner: web.AppRunner | None = None
        self._site: web.TCPSite | None = None

    @property
    def liveness_probe(self) -> LivenessProbe:
        """Get liveness probe for adding custom checks."""
        return self._liveness

    @property
    def readiness_probe(self) -> ReadinessProbe:
        """Get readiness probe for adding custom checks."""
        return self._readiness

    def mark_ready(self) -> None:
        """Mark the application as ready to receive traffic."""
        self._readiness.mark_ready()

    def mark_not_ready(self) -> None:
        """Mark the application as not ready."""
        self._readiness.mark_not_ready()

    def set_state_getter(self, getter: StateGetter) -> None:
        """Set the callback for getting processor state."""
        self._state_getter = getter

    def set_state_restorer(self, restorer: StateRestorer) -> None:
        """Set the callback for restoring processor state."""
        self._state_restorer = restorer

    def set_prestop_handler(self, handler: PreStopHandler) -> None:
        """Set the callback for PreStop hook."""
        self._prestop_handler = handler

    async def start(self) -> None:
        """
        Start the health server.

        Raises:
            DoryHealthError: If server fails to start
        """
        try:
            self._app = web.Application()
            self._setup_routes()

            self._runner = web.AppRunner(self._app)
            await self._runner.setup()

            self._site = web.TCPSite(
                self._runner,
                host="0.0.0.0",
                port=self._port,
            )
            await self._site.start()

            logger.info(f"Health server started on port {self._port}")

        except Exception as e:
            raise DoryHealthError(f"Failed to start health server: {e}", cause=e)

    async def stop(self) -> None:
        """Stop the health server."""
        if self._runner:
            await self._runner.cleanup()
            logger.info("Health server stopped")

    def _setup_routes(self) -> None:
        """Configure HTTP routes."""
        self._app.router.add_get(self._health_path, self._handle_health)
        self._app.router.add_get(self._ready_path, self._handle_ready)
        self._app.router.add_get(self._metrics_path, self._handle_metrics)

        # State transfer endpoints (required by Dory Orchestrator)
        self._app.router.add_get("/state", self._handle_state_get)
        self._app.router.add_post("/state", self._handle_state_post)

        # PreStop hook endpoint (required by Dory Orchestrator)
        self._app.router.add_get("/prestop", self._handle_prestop)

        # Root endpoint for basic info
        self._app.router.add_get("/", self._handle_root)

    async def _handle_root(self, request: web.Request) -> web.Response:
        """Handle root endpoint."""
        return web.json_response({
            "service": "dory-processor",
            "endpoints": [
                self._health_path,
                self._ready_path,
                self._metrics_path,
                "/state",
                "/prestop",
            ],
        })

    async def _handle_health(self, request: web.Request) -> web.Response:
        """
        Handle liveness probe.

        Returns 200 if alive, 503 if unhealthy.
        """
        result = await self._liveness.check()

        status = 200 if result.healthy else 503
        return web.json_response(result.to_dict(), status=status)

    async def _handle_ready(self, request: web.Request) -> web.Response:
        """
        Handle readiness probe.

        Returns 200 if ready, 503 if not ready.
        """
        result = await self._readiness.check()

        status = 200 if result.healthy else 503
        return web.json_response(result.to_dict(), status=status)

    async def _handle_metrics(self, request: web.Request) -> web.Response:
        """
        Handle Prometheus metrics endpoint.

        Returns metrics in Prometheus text format.
        """
        if self._metrics_collector is None:
            return web.Response(
                text="# No metrics collector configured\n",
                content_type="text/plain",
            )

        try:
            metrics_text = self._metrics_collector.export_prometheus()
            return web.Response(
                text=metrics_text,
                content_type="text/plain; version=0.0.4",
                charset="utf-8",
            )
        except Exception as e:
            logger.error(f"Error exporting metrics: {e}")
            return web.Response(
                text=f"# Error exporting metrics: {e}\n",
                content_type="text/plain",
                status=500,
            )

    async def _handle_state_get(self, request: web.Request) -> web.Response:
        """
        Handle GET /state for state capture during migration.

        Called by Dory Orchestrator to capture state from old pod
        before transferring to new pod.

        Returns:
            JSON response with processor state
        """
        if self._state_getter is None:
            logger.warning("State getter not configured, returning empty state")
            return web.json_response({
                "error": "state_getter not configured",
                "data": {},
            }, status=503)

        try:
            state = self._state_getter()
            logger.info("State captured for transfer", extra={"state_keys": list(state.keys())})
            return web.json_response(state)
        except Exception as e:
            logger.error(f"Failed to capture state: {e}")
            return web.json_response(
                {"error": f"Failed to capture state: {e}"},
                status=500,
            )

    async def _handle_state_post(self, request: web.Request) -> web.Response:
        """
        Handle POST /state for state restoration during migration.

        Called by Dory Orchestrator to restore state to new pod
        after capturing from old pod.

        Returns:
            JSON response confirming state restoration
        """
        if self._state_restorer is None:
            logger.warning("State restorer not configured")
            return web.json_response({
                "error": "state_restorer not configured",
            }, status=503)

        try:
            state = await request.json()
            logger.info("Restoring state from transfer", extra={"state_keys": list(state.keys())})
            await self._state_restorer(state)
            logger.info("State restored successfully")
            return web.json_response({"status": "ok", "message": "State restored"})
        except Exception as e:
            logger.error(f"Failed to restore state: {e}")
            return web.json_response(
                {"error": f"Failed to restore state: {e}"},
                status=500,
            )

    async def _handle_prestop(self, request: web.Request) -> web.Response:
        """
        Handle GET /prestop for PreStop hook.

        Called by Kubernetes PreStop hook before pod termination.
        Allows the application to prepare for graceful shutdown.

        Returns:
            JSON response confirming prestop handling
        """
        logger.info("PreStop hook invoked - preparing for shutdown")

        # Mark as not ready to stop receiving new traffic
        self._readiness.mark_not_ready()

        if self._prestop_handler:
            try:
                await self._prestop_handler()
                logger.info("PreStop handler completed")
            except Exception as e:
                logger.error(f"PreStop handler error: {e}")
                # Continue anyway - don't block shutdown

        return web.json_response({
            "status": "ok",
            "message": "PreStop hook processed, ready for termination",
        })
