"""Edge hardening: run throttling, security headers.

Body caps come from Starlette's RequestBodyLimitMiddleware (wired in
main.py). All limits are settings-driven so tests can shrink them. State
is in-process memory: correct for the single-replica SQLite deployment;
revisit with Redis if the API ever scales horizontally.
"""

from __future__ import annotations

import time
from collections import defaultdict

from starlette.datastructures import MutableHeaders
from starlette.requests import Request
from starlette.responses import JSONResponse
from starlette.types import ASGIApp, Message, Receive, Scope, Send

from app.settings import settings

# POST .../run is the expensive endpoint (LFX execution, LLM cost):
# per client IP and minute.
RUN_PATH_SUFFIX = "/run"
THROTTLE_WINDOW_SECONDS = 60.0

_hits: dict[str, list[float]] = defaultdict(list)


def reset_throttle() -> None:
    """Drop all counters (tests)."""
    _hits.clear()


def _client_ip(request: Request) -> str:
    # Behind the TLS proxy ProxyHeadersMiddleware has already rewritten
    # scope["client"] from X-Forwarded-For.
    if request.client is not None:
        return request.client.host
    forwarded = request.headers.get("x-forwarded-for", "")
    return forwarded.split(",")[0].strip() or "unknown"


class RunThrottleMiddleware:
    """Per-IP sliding window on POST .../run (429 + Retry-After)."""

    def __init__(self, app: ASGIApp) -> None:
        self.app = app

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        if (
            scope["type"] != "http"
            or scope.get("method") != "POST"
            or not scope.get("path", "").endswith(RUN_PATH_SUFFIX)
        ):
            await self.app(scope, receive, send)
            return
        request = Request(scope)
        now = time.monotonic()
        window_start = now - THROTTLE_WINDOW_SECONDS
        bucket = [t for t in _hits[_client_ip(request)] if t > window_start]
        if len(bucket) >= settings.run_throttle_per_minute:
            retry_after = max(1, int(bucket[0] + THROTTLE_WINDOW_SECONDS - now))
            response = JSONResponse({"detail": "too many runs, try again later"}, status_code=429)
            response.headers["Retry-After"] = str(retry_after)
            await response(scope, receive, send)
            return
        bucket.append(now)
        _hits[_client_ip(request)] = bucket
        await self.app(scope, receive, send)


class SecurityHeadersMiddleware:
    """Baseline response headers (TLS/HSTS live on the proxy)."""

    def __init__(self, app: ASGIApp) -> None:
        self.app = app

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

        async def secured(message: Message) -> None:
            if message["type"] == "http.response.start":
                headers = MutableHeaders(raw=message.setdefault("headers", []))
                headers["X-Content-Type-Options"] = "nosniff"
                headers["X-Frame-Options"] = "SAMEORIGIN"
                headers["Referrer-Policy"] = "strict-origin-when-cross-origin"
                headers["Content-Security-Policy"] = "frame-ancestors 'self'"
            await send(message)

        await self.app(scope, receive, secured)