"""SP19 — Security middleware + health probe.

Three concrete middlewares (body size, rate limit, security headers)
plus a richer ``/api/health`` snapshot. Sizing is for Cyclone's
local-only posture: a misconfigured Tailscale / ngrok bind, a
misbehaving cron job, a port-scanner scraping the API. Anything more
aggressive (auth, mTLS, WAF) is out of scope.

Design choices
--------------

* **In-memory rate limiter.** Cyclone is single-process; a dict
  keyed by IP is enough. If we ever go multi-worker, swap for
  Redis. The rate-limit counter resets after the bucket window;
  failing open on the limiter itself (an unexpected exception)
  rather than 503ing every request is the right call for a local tool.

* **Body-size check by Content-Length first, then chunked-read
  guard.** A chunked POST can lie about its size (or omit the
  header entirely); we cap read body size on the underlying stream
  so a malicious client can't keep streaming forever.

* **Security headers on every response.** CSP locks the API to
  same-origin + the Vite dev origin (whitelisted explicitly so a
  future operator running on a different port doesn't break).

* **Health snapshot is best-effort.** Each subsystem (DB,
  scheduler, pubsub) reports independently — a DB outage doesn't
  blank out the rest. ``status: "degraded"`` if any subsystem is
  unhappy; ``"ok"`` only when everything is.
"""
from __future__ import annotations

import json
import logging
import os
import threading
import time
from collections import deque
from dataclasses import dataclass, field
from typing import Any, Callable

from fastapi import Request, Response
from fastapi.responses import JSONResponse
from starlette.types import ASGIApp, Message, Receive, Scope, Send

log = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Knobs (env-var driven)
# ---------------------------------------------------------------------------


DEFAULT_MAX_BODY_BYTES = 50 * 1024 * 1024  # 50 MB — generous for X12 EDI
DEFAULT_RATE_LIMIT_PER_MIN = 300
DEFAULT_RATE_LIMIT_WINDOW_S = 60

# CSP: API responses are JSON, not HTML. ``default-src 'none'`` is the
# strictest setting; it forbids the API from being a vector for
# injected scripts in case an operator opens a JSON viewer with an
# HTML renderer.
_SECURITY_HEADERS: dict[str, str] = {
    "X-Content-Type-Options": "nosniff",
    "X-Frame-Options": "DENY",
    "Referrer-Policy": "same-origin",
    "Permissions-Policy": "geolocation=(), microphone=(), camera=()",
    "Content-Security-Policy": "default-src 'none'; frame-ancestors 'none'",
}


def _env_int(name: str, default: int) -> int:
    raw = os.environ.get(name)
    if not raw:
        return default
    try:
        return int(raw)
    except ValueError:
        log.warning("SP19: %s=%r is not an int; using default %d", name, raw, default)
        return default


# ---------------------------------------------------------------------------
# Body-size middleware
# ---------------------------------------------------------------------------


class BodySizeLimitMiddleware:
    """Reject requests whose body exceeds ``max_bytes``.

    Pure ASGI middleware (not BaseHTTPMiddleware — that one breaks
    FastAPI's ``request.body()`` introspection). Two-stage guard:
      1. If the request declares a ``Content-Length`` larger than
         ``max_bytes``, reject immediately with ``413``.
      2. While reading the body chunks, cap accumulated bytes at
         ``max_bytes``. If we cross the cap, return 413 instead of
         letting the handler read the rest.
    """

    def __init__(self, app: ASGIApp, max_bytes: int | None = None) -> None:
        self.app = app
        self.max_bytes = max_bytes or _env_int(
            "CYCLONE_MAX_BODY_BYTES", DEFAULT_MAX_BODY_BYTES,
        )

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        if scope["type"] != "http":
            await self.app(scope, receive, send)
            return

        # Stage 1: declared length.
        cl_header = None
        for k, v in scope.get("headers", []):
            if k == b"content-length":
                cl_header = v.decode("latin-1")
                break
        if cl_header is not None:
            try:
                if int(cl_header) > self.max_bytes:
                    await _send_rejection(
                        scope, send,
                        code=413,
                        reason="body_too_large",
                        detail=f"Content-Length {cl_header} exceeds limit {self.max_bytes}",
                    )
                    return
            except ValueError:
                await _send_rejection(
                    scope, send,
                    code=400, reason="bad_content_length",
                    detail=f"Content-Length {cl_header!r} is not an integer",
                )
                return

        # Stage 2: chunked read guard.
        seen = 0
        over_limit = False

        async def wrapped_receive() -> Message:
            nonlocal seen, over_limit
            if over_limit:
                # Drain any remaining bytes so the upstream ASGI
                # server doesn't see a truncated stream.
                msg = await receive()
                if msg.get("type") == "http.request":
                    return {"type": "http.request", "body": b"", "more_body": False}
                return msg
            msg = await receive()
            if msg.get("type") == "http.request":
                body = msg.get("body", b"") or b""
                seen += len(body)
                if seen > self.max_bytes:
                    over_limit = True
                    return {"type": "http.request", "body": b"", "more_body": False}
            return msg

        if cl_header is None:
            # Chunked / unknown length — guard with wrapped receive.
            await self.app(scope, wrapped_receive, send)
        else:
            # Fixed-length known to be safe; pass through.
            await self.app(scope, receive, send)


# ---------------------------------------------------------------------------
# Rate-limit middleware
# ---------------------------------------------------------------------------


@dataclass
class _Bucket:
    """Sliding-window counter for one IP."""
    timestamps: deque = field(default_factory=deque)

    def hit(self, window_s: int, now: float) -> bool:
        """Record one hit; return True if under the limit, False if over."""
        # Drop expired entries.
        cutoff = now - window_s
        while self.timestamps and self.timestamps[0] < cutoff:
            self.timestamps.popleft()
        return True  # we always record; the dispatcher decides to reject

    def count_in_window(self, now: float, window_s: int) -> int:
        cutoff = now - window_s
        while self.timestamps and self.timestamps[0] < cutoff:
            self.timestamps.popleft()
        return len(self.timestamps)


class RateLimitMiddleware:
    """Per-IP sliding-window rate limiter (pure ASGI).

    Defaults to ``CYCLONE_RATE_LIMIT_PER_MIN`` requests/minute per IP.
    Health-check probes and the ``/api/health`` endpoint are exempt
    so a load balancer's frequent probes don't trip the limiter.

    On unexpected errors the limiter fails OPEN — better to serve a
    few extra requests than to 503 every request because of a bug.
    """

    EXEMPT_PATHS = ("/api/health", "/healthz", "/readyz")

    def __init__(
        self,
        app: ASGIApp,
        per_minute: int | None = None,
        window_s: int | None = None,
    ) -> None:
        self.app = app
        self.per_minute = per_minute or _env_int(
            "CYCLONE_RATE_LIMIT_PER_MIN", DEFAULT_RATE_LIMIT_PER_MIN,
        )
        self.window_s = window_s or DEFAULT_RATE_LIMIT_WINDOW_S
        self._buckets: dict[str, _Bucket] = {}
        self._lock = threading.Lock()

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        if scope["type"] != "http":
            await self.app(scope, receive, send)
            return
        path = scope.get("path", "")
        if path in self.EXEMPT_PATHS:
            await self.app(scope, receive, send)
            return
        ip = _client_ip_from_scope(scope)
        now = time.monotonic()
        try:
            with self._lock:
                bucket = self._buckets.setdefault(ip, _Bucket())
                bucket.timestamps.append(now)
                count = bucket.count_in_window(now, self.window_s)
            if count > self.per_minute:
                await _send_rejection(
                    scope, send,
                    code=429,
                    reason="rate_limited",
                    detail=(
                        f"IP {ip} exceeded {self.per_minute} req/"
                        f"{self.window_s}s window"
                    ),
                )
                return
        except Exception as exc:  # noqa: BLE001
            log.warning("SP19: rate limiter failed open: %s", exc)
        await self.app(scope, receive, send)


def _client_ip_from_scope(scope: Scope) -> str:
    """Best-effort client IP from the ASGI scope. Falls back to ``"unknown"``."""
    for k, v in scope.get("headers", []):
        if k == b"x-forwarded-for":
            return v.decode("latin-1").split(",")[0].strip()
    client = scope.get("client")
    if client and client[0]:
        return client[0]
    return "unknown"


def _client_ip(request: Request) -> str:
    """Legacy helper (kept for the audit-event log path)."""
    return _client_ip_from_scope(request.scope)


# ---------------------------------------------------------------------------
# Security-headers middleware
# ---------------------------------------------------------------------------


class SecurityHeadersMiddleware:
    """Stamp the static security headers on every response (pure ASGI).

    CSP / X-Content-Type-Options / X-Frame-Options / Referrer-Policy /
    Permissions-Policy. The headers are static for now; per-route
    overrides can be added later if a route needs to relax them.
    """

    def __init__(self, app: ASGIApp, extra: dict[str, str] | None = None) -> None:
        self.app = app
        self.headers = [(k.lower().encode("latin-1"), v.encode("latin-1"))
                        for k, v in _SECURITY_HEADERS.items()]
        if extra:
            self.headers.extend(
                (k.lower().encode("latin-1"), v.encode("latin-1"))
                for k, v in extra.items()
            )

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        if scope["type"] != "http":
            await self.app(scope, receive, send)
            return

        async def wrapped_send(message: Message) -> None:
            if message["type"] == "http.response.start":
                headers = list(message.get("headers", []))
                existing = {k for k, _ in headers}
                for k, v in self.headers:
                    if k not in existing:
                        headers.append((k, v))
                message["headers"] = headers
            await send(message)

        await self.app(scope, receive, wrapped_send)


# ---------------------------------------------------------------------------
# Reject helper (also writes an audit event when DB is available)
# ---------------------------------------------------------------------------


async def _send_rejection(
    scope: Scope,
    send: Send,
    *,
    code: int,
    reason: str,
    detail: str,
) -> None:
    """Build a 413/429 JSON response, send it, and emit a log + audit event."""
    method = scope.get("method", "GET")
    path = scope.get("path", "/")
    ip = _client_ip_from_scope(scope)
    log.warning(
        "api.request_rejected",
        extra={
            "status": code,
            "reason": reason,
            "path": path,
            "method": method,
            "ip": ip,
            "detail": detail,
        },
    )
    payload = {"error": reason, "detail": detail, "status": code}
    body = json.dumps(payload).encode("utf-8")
    await send({
        "type": "http.response.start",
        "status": code,
        "headers": [
            (b"content-type", b"application/json"),
            (b"content-length", str(len(body)).encode("latin-1")),
        ],
    })
    await send({"type": "http.response.body", "body": body, "more_body": False})
    # Best-effort audit-log append. Don't block the response on a DB
    # outage (the rejection is the more important signal anyway).
    try:
        from cyclone import db
        from cyclone.audit_log import AuditEvent, append_event

        with db.SessionLocal()() as session:
            append_event(
                session,
                AuditEvent(
                    event_type="api.request_rejected",
                    entity_type="http_request",
                    entity_id=f"{method} {path}",
                    payload={
                        "status": code,
                        "reason": reason,
                        "path": path,
                        "method": method,
                        "ip": ip,
                    },
                    actor=f"api:{ip}",
                ),
            )
            session.commit()
    except Exception as exc:  # noqa: BLE001
        log.debug("SP19: audit-log append failed for rejection: %s", exc)


def _reject(
    request: Request,
    *,
    code: int,
    reason: str,
    detail: str,
) -> JSONResponse:
    """Sync helper kept for back-compat (the ``audit_log`` payload path)."""
    return JSONResponse(
        {"error": reason, "detail": detail, "status": code},
        status_code=code,
    )


# ---------------------------------------------------------------------------
# Health snapshot
# ---------------------------------------------------------------------------


@dataclass
class HealthSnapshot:
    status: str
    version: str
    db: dict[str, Any] = field(default_factory=dict)
    scheduler: dict[str, Any] = field(default_factory=dict)
    pubsub: dict[str, Any] = field(default_factory=dict)
    batch: dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> dict[str, Any]:
        return {
            "status": self.status,
            "version": self.version,
            "db": self.db,
            "scheduler": self.scheduler,
            "pubsub": self.pubsub,
            "batch": self.batch,
        }


def get_health_snapshot() -> HealthSnapshot:
    """Gather a best-effort snapshot of every Cyclone subsystem.

    Returns ``HealthSnapshot`` with ``status="ok"`` only if every
    subsystem check passes. ``"degraded"`` if any subsystem is
    unhappy but the API itself is responsive. Each subsystem reports
    independently so one outage doesn't blank out the rest.
    """
    from cyclone import __version__, db

    snap = HealthSnapshot(status="ok", version=__version__)

    # DB connectivity.
    try:
        with db.SessionLocal()() as session:
            session.execute(db.text("SELECT 1"))
        snap.db = {"ok": True}
    except Exception as exc:  # noqa: BLE001
        snap.db = {"ok": False, "error": str(exc)}
        snap.status = "degraded"

    # Scheduler state.
    try:
        from cyclone import scheduler as scheduler_mod
        sched = scheduler_mod.get_scheduler()
        snap.scheduler = {
            "running": sched.is_running(),
            "interval_s": sched._poll_interval,  # noqa: SLF001
            "sftp_block": sched._sftp_block_name,  # noqa: SLF001
        }
    except RuntimeError:
        snap.scheduler = {"running": False, "configured": False}
    except Exception as exc:  # noqa: BLE001
        snap.scheduler = {"ok": False, "error": str(exc)}
        snap.status = "degraded"

    # Backup scheduler.
    try:
        from cyclone import backup_scheduler as bks_mod
        bks = bks_mod.get_backup_scheduler()
        snap.scheduler["backup_scheduler_running"] = bks.is_running()
        snap.scheduler["backup_interval_hours"] = bks.interval_hours
    except (RuntimeError, ImportError):
        snap.scheduler["backup_scheduler_running"] = False
    except Exception:  # noqa: BLE001
        pass  # secondary subsystem; don't degrade the overall status

    # Pubsub bus stats — placeholder. The /api/health handler fills
    # in the real subscriber counts using request.app.state.event_bus.
    snap.pubsub = {"note": "filled in by health router"}

    # Last batch timestamp + count.
    try:
        from cyclone import db as db_mod
        from cyclone.db import Batch
        with db_mod.SessionLocal()() as session:
            row = (
                session.query(Batch)
                .order_by(Batch.parsed_at.desc())
                .first()
            )
        if row is not None:
            snap.batch = {
                "last_batch_id": row.id,
                "last_batch_kind": row.kind,
                "last_batch_at": row.parsed_at.isoformat() if row.parsed_at else None,
                "last_batch_filename": row.input_filename,
            }
        else:
            snap.batch = {"last_batch_id": None, "note": "no batches yet"}
    except Exception as exc:  # noqa: BLE001
        snap.batch = {"ok": False, "error": str(exc)}
        snap.status = "degraded"

    return snap
