"""HTTP middleware for Albus server.

Provides:
- Request ID tracking for correlation
- Human-friendly access logging
- Error handling with sanitization
- Response time tracking
- Configurable CORS
- Request timeouts
- API key authentication
- Rate limiting (per-IP, sliding window)
- Request body size enforcement
"""

from __future__ import annotations

import asyncio
import logging
import time
import uuid
from collections import defaultdict
from collections.abc import Awaitable, Callable
from typing import Any

from aiohttp import web

from core.observability import reset_log_request_id, set_log_request_id
from server.errors import ErrorCode, create_error_response, sanitize_error_message

# aiohttp middleware handler signature
_Handler = Callable[[web.Request], Awaitable[web.Response]]
# aiohttp middleware factory signature
_Middleware = Callable[[web.Request, _Handler], Awaitable[web.Response]]

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Shared state for rate limiting (per-worker; fine for single-process deploy)
# ---------------------------------------------------------------------------

_rate_buckets: dict[str, list[float]] = defaultdict(list)
"""IP -> list of request timestamps in the current window."""


# ---------------------------------------------------------------------------
# Request ID
# ---------------------------------------------------------------------------


@web.middleware
async def request_id_middleware(request: web.Request, handler: _Handler) -> web.Response:
    """Add request ID to all requests and set up logging context.

    This middleware:
    - Generates or extracts request ID from headers
    - Sets up observability context for logging
    - Adds request ID to response headers for correlation
    """
    request_id = request.headers.get("X-Request-ID") or f"req_{uuid.uuid4().hex[:12]}"
    request["request_id"] = request_id

    tok = set_log_request_id(request_id)
    try:
        response = await handler(request)
        response.headers["X-Request-ID"] = request_id
        return response
    finally:
        reset_log_request_id(tok)


# ---------------------------------------------------------------------------
# Access logging
# ---------------------------------------------------------------------------

# Paths exempt from access logging to reduce noise.
_ACCESS_LOG_SKIP_PATHS: frozenset[str] = frozenset(
    {"/health", "/healthz", "/readyz", "/api/v1/health", "/api/v1/healthz",
     "/api/v1/readyz", "/", "/api/v1", "/docs", "/openapi.json"}
)


@web.middleware
async def access_log_middleware(request: web.Request, handler: _Handler) -> web.Response:
    """Log HTTP requests with human-friendly format.

    Logs:
    - Method and path
    - Response status
    - Duration
    - Request ID for correlation
    """
    request_id = request.get("request_id", "unknown")
    start_time = time.time()
    method = request.method
    path = request.path_qs

    should_log = path not in _ACCESS_LOG_SKIP_PATHS

    # Get trace context if available
    trace_id = request.get("trace_id")
    span_id = request.get("span_id")

    if should_log:
        log_extra: dict[str, Any] = {"request_id": request_id, "method": method, "path": path}
        if trace_id:
            log_extra["trace_id"] = trace_id
        if span_id:
            log_extra["span_id"] = span_id

        trace_info = ""
        if trace_id:
            trace_info = f"  [trace={trace_id[:16]}]"

        logger.info(
            "→ %s %s  [request_id=%s]%s",
            method,
            path,
            request_id,
            trace_info,
            extra=log_extra,
        )

    try:
        response = await handler(request)
        duration_ms = (time.time() - start_time) * 1000
        status = response.status

        if should_log:
            status_emoji = "✓" if 200 <= status < 300 else "✗" if status >= 400 else "→"
            log_extra = {
                "request_id": request_id,
                "method": method,
                "path": path,
                "status": status,
                "duration_ms": duration_ms,
            }
            if trace_id:
                log_extra["trace_id"] = trace_id
            if span_id:
                log_extra["span_id"] = span_id

            trace_info = ""
            if trace_id:
                trace_info = f"  [trace={trace_id[:16]}]"

            logger.info(
                "%s %s %s  %d  %.1fms  [request_id=%s]%s",
                status_emoji,
                method,
                path,
                status,
                duration_ms,
                request_id,
                trace_info,
                extra=log_extra,
            )

        response.headers["X-Response-Time-Ms"] = f"{duration_ms:.2f}"
        return response
    except Exception as e:
        duration_ms = (time.time() - start_time) * 1000
        if should_log:
            logger.error(
                "✗ %s %s  ERROR  %.1fms  [request_id=%s]  %s",
                method,
                path,
                duration_ms,
                request_id,
                str(e),
                extra={
                    "request_id": request_id,
                    "method": method,
                    "path": path,
                    "duration_ms": duration_ms,
                    "error": str(e),
                },
                exc_info=True,
            )
        raise


# ---------------------------------------------------------------------------
# Error handler
# ---------------------------------------------------------------------------


@web.middleware
async def error_handler_middleware(request: web.Request, handler: _Handler) -> web.Response:
    """Standardize error responses with human-friendly messages.

    This middleware:
    - Catches all exceptions and converts to standardized error responses
    - Sanitizes error messages for production
    - Adds request ID and timestamp to errors
    - Logs errors with context
    """
    request_id = request.get("request_id", "unknown")
    start_time = time.time()

    try:
        response = await handler(request)
        response.headers["X-Response-Time-Ms"] = f"{(time.time() - start_time) * 1000:.2f}"
        return response
    except web.HTTPException as e:
        error_code = {
            404: ErrorCode.NOT_FOUND,
            401: ErrorCode.UNAUTHORIZED,
            403: ErrorCode.FORBIDDEN,
            429: ErrorCode.RATE_LIMIT_EXCEEDED,
        }.get(e.status, ErrorCode.BAD_REQUEST)

        error_msg = str(e.reason) or "Request failed"
        sanitized_msg = sanitize_error_message(error_msg)

        logger.warning(
            "HTTP %d: %s  [request_id=%s]",
            e.status,
            sanitized_msg,
            request_id,
            extra={"request_id": request_id, "status": e.status, "error": sanitized_msg},
        )

        return web.json_response(
            create_error_response(
                sanitized_msg,
                error_code,
                request_id=request_id,
            ),
            status=e.status,
        )
    except TimeoutError:
        duration_ms = (time.time() - start_time) * 1000
        logger.error(
            "Request timed out: %s %s  %.1fms  [request_id=%s]",
            request.method,
            request.path,
            duration_ms,
            request_id,
        )
        return web.json_response(
            create_error_response(
                "Request timed out",
                ErrorCode.TIMEOUT,
                request_id=request_id,
                details={"duration_ms": duration_ms},
            ),
            status=504,
        )
    except Exception as e:
        duration_ms = (time.time() - start_time) * 1000
        error_msg = str(e)
        sanitized_msg = sanitize_error_message(error_msg)

        logger.error(
            "Unhandled exception: %s  [request_id=%s]  %.1fms",
            sanitized_msg,
            request_id,
            duration_ms,
            extra={
                "request_id": request_id,
                "error": sanitized_msg,
                "duration_ms": duration_ms,
            },
            exc_info=True,
        )

        return web.json_response(
            create_error_response(
                sanitized_msg,
                ErrorCode.INTERNAL_ERROR,
                request_id=request_id,
                details={"duration_ms": duration_ms} if duration_ms > 1000 else None,
            ),
            status=500,
        )


# ---------------------------------------------------------------------------
# CORS (configurable)
# ---------------------------------------------------------------------------


def _make_cors_middleware(
    allowed_origins: list[str],
    allowed_methods: str,
    allowed_headers: str,
) -> _Middleware:
    """Create a CORS middleware with the given allowed origins.

    If allowed_origins contains '*', all origins are permitted (dev mode).
    Otherwise, the Origin header is checked against the allow-list.
    """
    allow_all = "*" in allowed_origins
    origin_set = frozenset(allowed_origins)

    @web.middleware
    async def cors_middleware(request: web.Request, handler: _Handler) -> web.Response:
        origin = request.headers.get("Origin", "")

        def _cors_headers() -> dict[str, str]:
            effective_origin = "*" if allow_all else (origin if origin in origin_set else "")
            if not effective_origin:
                return {}
            return {
                "Access-Control-Allow-Origin": effective_origin,
                "Access-Control-Allow-Methods": allowed_methods,
                "Access-Control-Allow-Headers": allowed_headers,
                "Access-Control-Max-Age": "86400",
            }

        if request.method == "OPTIONS":
            return web.Response(headers=_cors_headers())

        response = await handler(request)
        cors = _cors_headers()
        for key, value in cors.items():
            response.headers[key] = value
        return response

    return cors_middleware


# ---------------------------------------------------------------------------
# Request timeout
# ---------------------------------------------------------------------------


def _make_timeout_middleware(timeout_seconds: int) -> _Middleware:
    """Wrap every request in an asyncio timeout.

    On expiry the handler's task is cancelled and a 504 bubbles up through
    ``error_handler_middleware`` (which catches ``asyncio.TimeoutError``).
    """

    @web.middleware
    async def timeout_middleware(request: web.Request, handler: _Handler) -> web.Response:
        return await asyncio.wait_for(handler(request), timeout=timeout_seconds)

    return timeout_middleware


# ---------------------------------------------------------------------------
# API key authentication
# ---------------------------------------------------------------------------

# Paths that are exempt from API key checks (health, liveness, readiness, docs).
_AUTH_EXEMPT_PATHS: frozenset[str] = frozenset(
    {"/", "/api/v1", "/health", "/healthz", "/readyz",
     "/api/v1/health", "/api/v1/healthz", "/api/v1/readyz",
     "/docs", "/api/v1/docs", "/openapi.json", "/api/v1/openapi.json"}
)


def _make_api_key_middleware(valid_keys: list[str]) -> _Middleware:
    """Require ``X-API-Key`` header on every non-exempt request.

    Returns 401 with a JSON body if the key is missing or invalid.
    """
    key_set = frozenset(valid_keys)

    @web.middleware
    async def api_key_middleware(request: web.Request, handler: _Handler) -> web.Response:
        if request.path in _AUTH_EXEMPT_PATHS or request.method == "OPTIONS":
            return await handler(request)

        supplied = request.headers.get("X-API-Key", "")
        if supplied not in key_set:
            raise web.HTTPUnauthorized(reason="Invalid or missing API key")

        return await handler(request)

    return api_key_middleware


# ---------------------------------------------------------------------------
# Rate limiting (sliding-window per IP)
# ---------------------------------------------------------------------------


def _make_rate_limit_middleware(rpm: int) -> _Middleware:
    """Enforce a per-IP requests-per-minute limit.

    Uses a simple sliding-window counter stored in ``_rate_buckets``.
    Returns 429 when the limit is exceeded, with a ``Retry-After`` header.
    """
    window = 60.0  # 1 minute

    @web.middleware
    async def rate_limit_middleware(request: web.Request, handler: _Handler) -> web.Response:
        # Health endpoints are exempt from rate limiting
        if request.path in _AUTH_EXEMPT_PATHS:
            return await handler(request)

        client_ip = request.remote or "unknown"
        now = time.time()
        cutoff = now - window

        # Prune old entries and append current
        bucket = _rate_buckets[client_ip]
        _rate_buckets[client_ip] = [ts for ts in bucket if ts > cutoff]
        _rate_buckets[client_ip].append(now)

        if len(_rate_buckets[client_ip]) > rpm:
            retry_after = int(window - (now - _rate_buckets[client_ip][0])) + 1
            raise web.HTTPTooManyRequests(
                reason="Rate limit exceeded",
                headers={"Retry-After": str(max(1, retry_after))},
            )

        return await handler(request)

    return rate_limit_middleware


# ---------------------------------------------------------------------------
# Public factory — builds the full middleware stack from config
# ---------------------------------------------------------------------------


def build_middleware_stack(
    *,
    request_timeout: int = 300,
    cors_allowed_origins: list[str] | None = None,
    cors_allowed_methods: str = "GET, POST, DELETE, PATCH, OPTIONS",
    cors_allowed_headers: str = "Content-Type, Authorization, X-Request-ID, X-API-Key",
    api_keys: list[str] | None = None,
    rate_limit_rpm: int = 0,
) -> list[_Middleware]:
    """Assemble the ordered middleware stack from config values.

    Order (outermost → innermost):
        tracing → request_id → access_log → error_handler → timeout
        → rate_limit (opt) → api_key (opt) → cors
    """
    from server.transport.tracing import tracing_middleware

    stack: list[_Middleware] = [
        tracing_middleware,
        request_id_middleware,
        access_log_middleware,
        error_handler_middleware,
        _make_timeout_middleware(request_timeout),
    ]

    if rate_limit_rpm > 0:
        stack.append(_make_rate_limit_middleware(rate_limit_rpm))

    if api_keys:
        stack.append(_make_api_key_middleware(api_keys))

    stack.append(
        _make_cors_middleware(
            cors_allowed_origins or ["*"],
            cors_allowed_methods,
            cors_allowed_headers,
        )
    )

    return stack


__all__ = [
    "access_log_middleware",
    "build_middleware_stack",
    "error_handler_middleware",
    "request_id_middleware",
]
