"""Public listener: path allowlist, model resolution, streaming proxy.

Security shape (plan section 6):
  * the *normalized* path decides what is served, so `/v1/../sleep` and
    `/v1%2f..%2fsleep` both collapse to `/sleep` and get a 404;
  * `/admin/*` and every vLLM dev endpoint are simply not in the allowlist;
  * the admin listener lives on another socket (see admin_api.py).
"""

from __future__ import annotations

import json
import logging
import posixpath
import re
import time
import urllib.parse
from typing import Any

import httpx
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse, Response, StreamingResponse
from starlette.background import BackgroundTask
from starlette.datastructures import MutableHeaders
from starlette.middleware.body_limit import RequestBodyLimitMiddleware

from config import (
    DEPTH_AWAKE,
    DEPTH_OFFLOADED,
    DEPTH_RESTARTING,
    DEPTH_SLEEPING,
    Config,
)
from services import ServiceManager, ServiceState

log = logging.getLogger("vllm_router.http")

# Endpoints whose bodies legitimately omit `model` (plan 6.2 / 6.6).
DEFAULT_TEXT_PATHS = frozenset({"/v1/chat/completions", "/v1/completions"})
# A multipart request on the chat path is the image-bearing (OCR) path.
MULTIPART_DEFAULT_PATH = "/v1/chat/completions"

# Headers that must never be forwarded in either direction.
HOP_BY_HOP = frozenset({
    "connection", "keep-alive", "proxy-authenticate", "proxy-authorization",
    "te", "trailers", "transfer-encoding", "upgrade",
})
REQUEST_ONLY_STRIP = HOP_BY_HOP | {"host", "content-length", "expect"}
RESPONSE_STRIP = HOP_BY_HOP | {"content-length"}

_BOUNDARY_RE = re.compile(r'boundary="?([^";,]+)"?', re.IGNORECASE)
_NAME_RE = re.compile(r'name="((?:[^"\\]|\\.)*)"', re.IGNORECASE)
# A text field we are willing to read fully (and re-send) while sniffing.
_MAX_MODEL_FIELD_BYTES = 4096


# --------------------------------------------------------------------------
# errors (OpenAI-shaped so existing clients can parse them)
# --------------------------------------------------------------------------
def _error(status: int, message: str, *, err_type: str = "invalid_request_error",
           code: str | None = None, param: str | None = None,
           extra: dict[str, Any] | None = None, headers: dict[str, str] | None = None
           ) -> JSONResponse:
    err: dict[str, Any] = {"message": message, "type": err_type}
    if param is not None:
        err["param"] = param
    if code is not None:
        err["code"] = code
    if extra:
        err.update(extra)
    return JSONResponse({"error": err}, status_code=status, headers=headers)


def not_found_response(message: str = "Not found") -> JSONResponse:
    """Uniform 404 for everything outside the allowlist (paths *and* models)."""
    return _error(404, message, code="not_found")


def shutting_down_response(detail: str = "router is shutting down") -> JSONResponse:
    """503 for in-flight work that dies because the router is stopping."""
    return JSONResponse(
        {"error": {
            "type": "router_shutting_down",
            "code": "router_shutting_down",
            "message": f"Request aborted: {detail}; retry shortly",
            "retry_after_seconds": 5,
        }},
        status_code=503,
        headers={"Retry-After": "5", "connection": "close"},
    )


_SHUTDOWN_BODY = json.dumps({
    "error": {
        "type": "router_shutting_down",
        "code": "router_shutting_down",
        "message": "Request aborted: the router is shutting down; retry shortly",
        "retry_after_seconds": 5,
    }
}).encode()


class ShutdownGuardMiddleware:
    """Turns shutdown-induced cancellation of an in-flight request into a
    depth-aware 503 instead of uvicorn's bare 21-byte 500 (E2E case 12).

    uvicorn cancels pending ASGI tasks once `timeout_graceful_shutdown`
    expires; that CancelledError is a BaseException, so Starlette's
    ServerErrorMiddleware (and every FastAPI exception_handler) ignores it and
    uvicorn answers with a plain-text 500.  This is the only layer that can
    intercept it.  A response whose body has already started streaming cannot
    be un-sent: those connections are simply closed (client sees a truncated
    stream), which is the best available behaviour.
    """

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

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

        async def send_wrapper(message) -> None:
            nonlocal started
            if message["type"] == "http.response.start":
                started = True
            await send(message)

        try:
            await self.app(scope, receive, send_wrapper)
        except BaseException as exc:  # noqa: BLE001 - CancelledError is one
            if started or not self.manager.shutting_down:
                raise
            log.info("abort_shutdown method=%s path=%s error=%s",
                     scope.get("method"), scope.get("path"), type(exc).__name__)
            try:
                await send({"type": "http.response.start", "status": 503, "headers": [
                    (b"content-type", b"application/json"),
                    (b"content-length", str(len(_SHUTDOWN_BODY)).encode("latin-1")),
                    (b"retry-after", b"5"),
                    (b"connection", b"close"),
                ]})
                await send({"type": "http.response.body", "body": _SHUTDOWN_BODY})
            except Exception:  # pragma: no cover - client may be gone already
                pass


def model_not_found_response(model: str) -> JSONResponse:
    return _error(
        404,
        f"The model `{model}` does not exist or is not served by this endpoint.",
        param="model",
        code="model_not_found",
    )


def waking_response(cfg: Config, svc: ServiceState, depth: str,
                    detail: str = "") -> JSONResponse:
    """Depth-aware 503 (plan 6.2.1)."""
    policy = cfg.policy(depth)
    body = policy.error_body(svc.cfg.model)
    if detail:
        body["error"]["message"] += f" ({detail})"
    return JSONResponse(
        body,
        status_code=503,
        headers={
            "Retry-After": str(policy.retry_after_s),
            "x-vllm-router": "model-waking",
        },
    )


# --------------------------------------------------------------------------
# path handling
# --------------------------------------------------------------------------
def normalize_raw_path(raw: bytes | str) -> str:
    """Decode %xx and collapse `.`/`..` segments BEFORE any allowlist check."""
    if isinstance(raw, (bytes, bytearray)):
        raw = bytes(raw).decode("latin-1")
    path = urllib.parse.unquote(raw)
    path = posixpath.normpath(path)
    if not path.startswith("/"):
        path = "/" + path
    return path


def classify_public_path(norm: str, cfg: Config) -> str | None:
    """Returns the handler kind, or None for 404."""
    if norm == "/health":
        return "health"
    if norm == "/metrics" and cfg.metrics_enabled:
        return "metrics"
    if norm == "/v1/models":
        return "models"
    if norm.startswith("/v1/") and ".." not in norm.split("/") and "\\" not in norm:
        return "api"
    return None


# --------------------------------------------------------------------------
# model resolution (plan 6.2 step 1)
# --------------------------------------------------------------------------
def extract_multipart_model(body: bytes, content_type: str) -> str | None:
    """Pull a small `model` text field out of a multipart body.

    Deliberately hand-rolled: Starlette's form parser would consume (and
    buffer/spool) every file part.  We walk the boundary delimiters, look at
    part headers only, and never slice a file part's payload.
    """
    match = _BOUNDARY_RE.search(content_type)
    if match is None:
        return None
    delimiter = b"--" + match.group(1).encode("latin-1")
    pos = 0
    while True:
        start = body.find(delimiter, pos)
        if start < 0:
            return None
        line_end = body.find(b"\r\n", start + len(delimiter))
        if line_end < 0:
            return None
        next_delim = body.find(delimiter, line_end)
        if next_delim < 0:
            return None
        headers = body[line_end + 2:next_delim]
        head, sep, _ = headers.partition(b"\r\n\r\n")
        if not sep:
            pos = next_delim
            continue
        disposition = b""
        for line in head.split(b"\r\n"):
            if line.lower().startswith(b"content-disposition:"):
                disposition = line
                break
        if b"filename=" in disposition.lower():
            pos = next_delim  # file part: leave it untouched
            continue
        name_match = _NAME_RE.search(disposition.decode("latin-1", "replace"))
        if name_match is None or name_match.group(1).lower() != "model":
            pos = next_delim
            continue
        value = headers[len(head) + 4:].rstrip(b"\r\n")
        if len(value) > _MAX_MODEL_FIELD_BYTES:
            return None
        return value.decode("utf-8", "replace").strip() or None


def resolve_model_name(cfg: Config, name: str) -> str | None:
    """Registry lookup (case-insensitive over model names, aliases and keys)."""
    return cfg.index.get(name.strip().lower())


async def resolve_target(manager: ServiceManager, request: Request, norm_path: str
                         ) -> tuple[ServiceState | None, Response | None]:
    """Decide which service a request belongs to.  Returns (service, error);
    exactly one is None.  Never wakes anything when the model is unknown."""
    cfg = manager.cfg

    # 1. /v1/embeddings always goes to the embed service.
    if norm_path == "/v1/embeddings":
        return manager.services["embed"], None

    # 2. Model named in the path (GET /v1/models/{id}).
    if norm_path.startswith("/v1/models/") and request.method in ("GET", "HEAD"):
        key = resolve_model_name(cfg, urllib.parse.unquote(norm_path[len("/v1/models/"):]))
        if key is None:
            return None, model_not_found_response(norm_path[len("/v1/models/"):])
        return manager.services[key], None

    content_type = request.headers.get("content-type", "")
    is_json = content_type.split(";")[0].strip().lower() in ("application/json",) or \
        content_type.strip().lower().endswith("+json")
    is_multipart = content_type.split(";")[0].strip().lower().startswith("multipart/")

    if is_json:
        raw = await request.body()
        try:
            payload = json.loads(raw or b"{}")
        except (ValueError, UnicodeDecodeError):
            return None, _error(400, "Request body is not valid JSON.", code="invalid_json")
        if not isinstance(payload, dict):
            return None, _error(400, "Request body must be a JSON object.", code="invalid_json")
        model = payload.get("model")
        if isinstance(model, str) and model.strip():
            key = resolve_model_name(cfg, model)
            if key is None:
                return None, model_not_found_response(model)
            return manager.services[key], None
        if norm_path in DEFAULT_TEXT_PATHS:
            return manager.services["text"], None
        return None, _error(
            400,
            "Missing required parameter: 'model'.",
            param="model",
            code="missing_model",
        )

    if is_multipart:
        model = extract_multipart_model(await request.body(), content_type)
        if model is not None:
            key = resolve_model_name(cfg, model)
            if key is None:
                return None, model_not_found_response(model)
            return manager.services[key], None
        if norm_path == MULTIPART_DEFAULT_PATH:
            return manager.services["ocr"], None  # image-bearing path
        if norm_path in DEFAULT_TEXT_PATHS:
            return manager.services["text"], None
        return None, _error(
            400,
            "Missing required parameter: 'model'.",
            param="model",
            code="missing_model",
        )

    # Any other body shape (empty, text/plain, ...): only the two chat paths
    # have a defensible default.
    if norm_path in DEFAULT_TEXT_PATHS:
        return manager.services["text"], None
    return None, _error(
        400,
        "Missing required parameter: 'model'.",
        param="model",
        code="missing_model",
    )


# --------------------------------------------------------------------------
# proxy
# --------------------------------------------------------------------------
def _forwarded_request_headers(request: Request) -> list[tuple[str, str]]:
    return [
        (k, v) for k, v in request.headers.items()
        if k.lower() not in REQUEST_ONLY_STRIP
    ]


async def _open_upstream(manager: ServiceManager, request: Request, svc: ServiceState,
                         norm_path: str) -> httpx.Response:
    query = request.scope.get("query_string", b"").decode("latin-1")
    url = svc.cfg.base_url + norm_path + (f"?{query}" if query else "")
    upstream_req = manager.proxy_client.build_request(
        request.method,
        url,
        headers=_forwarded_request_headers(request),
        content=await request.body(),
    )
    # Stream: the response body is handed to the client chunk by chunk and is
    # never buffered here (SSE safe, 900s read timeout upstream).
    return await manager.proxy_client.send(upstream_req, stream=True)


async def _proxy_response(manager: ServiceManager, request: Request, svc: ServiceState,
                          norm_path: str) -> Response:
    upstream = await _open_upstream(manager, request, svc, norm_path)
    headers = MutableHeaders()
    for key, value in upstream.headers.items():
        if key.lower() not in RESPONSE_STRIP:
            headers.append(key, value)

    async def finish() -> None:
        # Runs when the last chunk has been sent (or the client hung up) --
        # this is what keeps the idle manager from sleeping mid-stream.
        try:
            await upstream.aclose()
        finally:
            manager.end_request(svc)

    return StreamingResponse(
        upstream.aiter_raw(),
        status_code=upstream.status_code,
        headers=headers,
        background=BackgroundTask(finish),
    )


async def forward(manager: ServiceManager, request: Request, svc: ServiceState,
                  norm_path: str) -> Response:
    """Ensure the target is awake, then proxy.  The request stays counted as
    active for the whole life of the response body."""
    cfg = manager.cfg
    manager.begin_request(svc)
    handed_off = False
    try:
        for attempt in (1, 2):
            outcome = await manager.ensure_awake(svc.cfg.key)
            if not outcome.ok:
                log.info("reject service=%s status=503 depth=%s reason=%s",
                         svc.cfg.key, outcome.depth, outcome.reason)
                detail = "" if outcome.reason == "hold_timeout" else outcome.detail
                return waking_response(cfg, svc, outcome.depth, detail)
            try:
                response = await _proxy_response(manager, request, svc, norm_path)
                handed_off = True
                return response
            except httpx.TransportError as exc:
                # Backend went away between the health check and the proxy call.
                svc.reachable = False
                svc.depth = None
                svc.state_checked_at = float("-inf")
                svc.last_error = f"proxy: {type(exc).__name__}"
                log.warning("proxy_transport_error service=%s attempt=%d error=%s",
                            svc.cfg.key, attempt, type(exc).__name__)
                if attempt == 2:
                    return waking_response(cfg, svc, DEPTH_RESTARTING,
                                           "backend connection failed")
    finally:
        if not handed_off:
            manager.end_request(svc)
    return waking_response(cfg, svc, DEPTH_RESTARTING, "unreachable")  # pragma: no cover


# --------------------------------------------------------------------------
# public endpoints
# --------------------------------------------------------------------------
def _service_summary(manager: ServiceManager, live: bool) -> dict[str, Any]:
    out: dict[str, Any] = {}
    for svc in manager.services.values():
        out[svc.cfg.key] = {
            "service": svc.cfg.service,
            "model": svc.cfg.model,
            "reachable": svc.reachable,
            "sleeping": None if svc.depth is None else svc.depth != DEPTH_AWAKE,
            "depth": svc.depth,
            "wake_in_progress": svc.wake_in_progress,
            "active_requests": svc.active_requests,
        }
    return out


def build_public_app(manager: ServiceManager) -> FastAPI:
    cfg = manager.cfg
    app = FastAPI(
        title="vllm-router",
        version="1.0.0",
        docs_url=None,
        redoc_url=None,
        openapi_url=None,
    )
    # FastAPI does not forward **extra to Starlette, so the request-body cap
    # (ROUTER_MAX_BODY_BYTES) is wired as middleware explicitly.
    app.add_middleware(RequestBodyLimitMiddleware, max_body_size=cfg.max_body_bytes)
    # Added last => outermost.  Must sit outside everything because the
    # exception it converts is a BaseException that Starlette's
    # ServerErrorMiddleware deliberately does not catch.
    app.add_middleware(ShutdownGuardMiddleware, manager=manager)

    @app.exception_handler(Exception)
    async def internal_error(_request: Request, exc: Exception) -> JSONResponse:
        """Anything unhandled becomes a parseable JSON error, never a bare
        21-byte "Internal Server Error".  During shutdown that is a 503 with
        Retry-After (E2E case 12)."""
        if manager.shutting_down:
            log.info("abort_shutdown method=%s path=%s error=%s",
                     _request.method, _request.url.path, type(exc).__name__)
            return shutting_down_response(f"{type(exc).__name__} during shutdown")
        log.exception("internal_error error=%s", type(exc).__name__)
        return _error(500, f"Unhandled router error: {type(exc).__name__}",
                      err_type="internal_error", code="internal_error",
                      headers={"Retry-After": "1"})

    @app.get("/health")
    async def health() -> JSONResponse:
        """Router liveness + cached per-service summary.  Always 200: a sleeping
        backend is normal, not an outage (docker healthchecks must not flap)."""
        return JSONResponse({
            "status": "ok",
            "router": {
                "uptime_s": round(time.monotonic() - manager.started_at, 1),
                "public_port": cfg.public_port,
                "admin_port": cfg.admin_port,
                "services": len(manager.services),
            },
            "services": _service_summary(manager, live=False),
        })

    @app.get("/v1/models")
    async def models() -> JSONResponse:
        return JSONResponse({
            "object": "list",
            "data": [
                {"id": svc.cfg.model, "object": "model", "created": 0,
                 "owned_by": "vllm-router"}
                for svc in manager.services.values()
            ],
        })

    # NOTE: /metrics is deliberately NOT registered as an explicit route -- the
    # catch-all below re-checks it against the allowlist so it 404s whenever
    # ROUTER_METRICS is off.
    async def metrics() -> Response:
        lines = [
            "# HELP vllm_router_service_depth 0=awake 1=sleeping 2=offloaded 3=restarting",
            "# TYPE vllm_router_service_depth gauge",
        ]
        rank = {DEPTH_AWAKE: 0, DEPTH_SLEEPING: 1, DEPTH_OFFLOADED: 2, DEPTH_RESTARTING: 3}
        for svc in manager.services.values():
            lines.append(f'vllm_router_service_depth{{service="{svc.cfg.key}"}} '
                         f'{rank.get(svc.depth, 2)}')
            lines.append(f'vllm_router_active_requests{{service="{svc.cfg.key}"}} '
                         f'{svc.active_requests}')
        return Response("\n".join(lines) + "\n", media_type="text/plain; version=0.0.4")

    async def gate(request: Request) -> Response:
        raw = request.scope.get("raw_path") or request.url.path.encode()
        norm = normalize_raw_path(raw)
        kind = classify_public_path(norm, cfg)
        if kind is None:
            log.info("reject_path method=%s path=%s", request.method, norm)
            return not_found_response()
        if kind == "health":
            return await health()
        if kind == "models":
            return await models()
        if kind == "metrics":
            return await metrics()

        svc, error = await resolve_target(manager, request, norm)
        if error is not None:
            log.info("reject_model method=%s path=%s status=%s",
                     request.method, norm, error.status_code)
            return error
        return await forward(manager, request, svc, norm)

    # Registered last: explicit routes above win, everything else lands here
    # and is re-checked against the allowlist on the *normalized* path.
    app.router.add_route(
        "/{path:path}",
        gate,
        methods=["GET", "POST", "PUT", "PATCH", "DELETE", "HEAD", "OPTIONS"],
    )
    return app