"""TypeSafe SDK wire gateway to distinct shallow/full vLLM scoring engines."""

from __future__ import annotations

import argparse
import asyncio
from contextlib import asynccontextmanager
from dataclasses import dataclass
import hashlib
import hmac
import json
import math
import os
from pathlib import Path
import re
import uuid

import httpx
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse

from contracts import (
    ServingError,
    answer,
    compile_question,
    prepare_request,
    probabilities,
)


def canonical_root(value):
    url = httpx.URL(value)
    path = url.path.rstrip("/").removesuffix("/v1").rstrip("/")
    return str(url.copy_with(path=path))


@dataclass(frozen=True)
class Settings:
    high_url: str
    low_url: str
    high_model: str
    low_model: str
    api_key: str
    model_name: str = "APUS-OpenJev-v1-9B"
    max_body_bytes: int = 1_048_576
    max_questions: int = 16
    max_tokens: int = 8191
    max_inflight: int = 32
    timeout: float = 120.0
    high_model_path: str | None = None
    low_model_path: str | None = None

    def __post_init__(self):
        for value in (self.high_url, self.low_url):
            url = httpx.URL(value)
            if (
                url.scheme not in ("http", "https")
                or not url.host
                or url.username
                or url.password
                or url.query
                or url.fragment
            ):
                raise ValueError(
                    "Backend URLs must be HTTP(S) without credentials/query"
                )
        if canonical_root(self.high_url) == canonical_root(self.low_url):
            raise ValueError(
                "Low and high must use distinct engine URLs; no silent full-depth fallback"
            )
        if not all((self.high_model, self.low_model, self.api_key, self.model_name)):
            raise ValueError(
                "Both model aliases and a nonempty local API key are required"
            )
        if self.high_model == self.low_model:
            raise ValueError("Low and high engine model aliases must be distinct")

    @classmethod
    def from_env(cls):
        return cls(
            high_url=os.environ["JEV_HIGH_URL"],
            low_url=os.environ["JEV_LOW_URL"],
            high_model=os.environ["JEV_HIGH_MODEL"],
            low_model=os.environ["JEV_LOW_MODEL"],
            api_key=os.environ["JEV_API_KEY"],
            high_model_path=os.environ["JEV_HIGH_MODEL_PATH"],
            low_model_path=os.environ["JEV_LOW_MODEL_PATH"],
            model_name=os.environ.get("JEV_MODEL_NAME", "APUS-OpenJev-v1-9B"),
        )


def unique_object(pairs):
    result = {}
    for key, value in pairs:
        if key in result:
            raise ValueError("Duplicate JSON keys")
        result[key] = value
    return result


def reject_constant(value):
    raise ValueError("Nonfinite JSON values are forbidden")


def parse_scores(data, tokens):
    try:
        choices = data["choices"]
        if len(choices) != 1 or len(choices[0]["logprobs"]["top_logprobs"]) != 1:
            raise ValueError
        values = choices[0]["logprobs"]["top_logprobs"][0]
        if not isinstance(values, dict):
            raise ValueError
        by_id = {}
        for key, value in values.items():
            if (
                not re.fullmatch(r"token_id:(0|[1-9]\d*)", key)
                or type(value) not in (float, int)
                or not math.isfinite(value)
                or value > 0
            ):
                raise ValueError
            by_id[int(key.split(":")[1])] = value
        return [by_id[token] for token in tokens]
    except (KeyError, TypeError, ValueError, IndexError):
        raise ServingError(
            502,
            "invalid_backend_scores",
            "Backend omitted valid candidate token logprobs",
        ) from None


def endpoint(url, path):
    root = url.rstrip("/")
    return root + (path.removeprefix("/v1") if root.endswith("/v1") else path)


def file_sha(path):
    return hashlib.sha256(Path(path).read_bytes()).hexdigest()


def verify_artifacts(settings):
    """Bind configured backend roots to real model configs and export identity."""
    if not settings.high_model_path or not settings.low_model_path:
        raise ValueError("Both JEV_HIGH_MODEL_PATH and JEV_LOW_MODEL_PATH are required")
    roots = {
        "high": Path(settings.high_model_path).resolve(),
        "low": Path(settings.low_model_path).resolve(),
    }
    for effort, depth in (("high", 32), ("low", 16)):
        root = roots[effort]
        cfg = json.loads((root / "config.json").read_text())
        text = cfg.get("text_config", cfg)
        if (
            cfg.get("model_type") not in ("qwen3_5", "qwen3_5_text")
            or text.get("num_hidden_layers") != depth
            or len(text.get("layer_types", [])) != depth
            or text.get("hidden_size") != 4096
        ):
            raise ValueError(
                "Backend artifact does not declare the required execution depth"
            )
        if (root / "adapter_config.json").exists() or (
            root / "EXPORT_INCOMPLETE"
        ).exists():
            raise ValueError("Backend artifact is not a completed merged model")
        if not (root / "model.safetensors.index.json").is_file():
            raise ValueError("Backend artifact has no merged weight index")
    low, high = roots["low"], roots["high"]
    proof = json.loads((low / "export-low-provenance.json").read_text())
    if (
        proof.get("status") != "verified"
        or proof.get("actual_depth") != 16
        or proof.get("all_retained_tensors_bit_identical") is not True
    ):
        raise ValueError("Low artifact has not passed prefix export verification")
    for key, path in {
        "export_config_sha256": low / "config.json",
        "export_index_sha256": low / "model.safetensors.index.json",
        "source_config_sha256": high / "config.json",
        "source_index_sha256": high / "model.safetensors.index.json",
        "source_release_manifest_sha256": high / "release-manifest.json",
    }.items():
        if proof.get(key) != file_sha(path):
            raise ValueError(
                "Low export provenance does not match the configured source or export metadata"
            )
    return roots


def create_app(settings, tokenizer, transport=None, check_backends=True):
    @asynccontextmanager
    async def lifespan(app):
        async with httpx.AsyncClient(
            timeout=settings.timeout,
            trust_env=False,
            follow_redirects=False,
            transport=transport,
        ) as client:
            app.state.client = client
            if check_backends:
                roots = verify_artifacts(settings)
                for effort, url, alias in (
                    ("low", settings.low_url, settings.low_model),
                    ("high", settings.high_url, settings.high_model),
                ):
                    response = await client.get(endpoint(url, "/v1/models"))
                    response.raise_for_status()
                    matched = [m for m in response.json()["data"] if m["id"] == alias]
                    if (
                        len(matched) != 1
                        or not matched[0].get("root")
                        or Path(matched[0]["root"]).resolve() != roots[effort]
                    ):
                        raise RuntimeError(
                            "Engine alias/root does not match the verified depth artifact"
                        )
            app.state.ready = True
            yield
            app.state.ready = False

    app = FastAPI(lifespan=lifespan)
    app.state.ready = False
    app.state.inflight = 0

    def authorize(request):
        authorization = request.headers.get("authorization", "")
        if not hmac.compare_digest(
            authorization.encode(), ("Bearer " + settings.api_key).encode()
        ):
            raise ServingError(401, "unauthorized", "Invalid local API key")

    @app.middleware("http")
    async def request_id_header(request, call_next):
        response = await call_next(request)
        response.headers.setdefault("x-typesafe-request-id", str(uuid.uuid4()))
        return response

    @app.exception_handler(ServingError)
    async def serving_error(_request, exc):
        return JSONResponse(
            status_code=exc.status,
            content={"error": {"code": exc.code, "message": exc.message}},
        )

    @app.get("/health")
    async def health():
        return JSONResponse(
            status_code=200 if app.state.ready else 503,
            content={
                "ready": app.state.ready,
                "model": settings.model_name,
                "efforts": {"low": 16, "high": 32},
                "routing": "separate-vllm-engines",
                "automatic_routing": False,
            },
        )

    @app.get("/v1/models")
    async def models(request: Request):
        authorize(request)
        return {
            "models": [
                {"name": name, "description": description, "release_date": "2026-09-21"}
                for name, description in [
                    (
                        settings.model_name,
                        "APUS local decision model; explicit low16/high32; uncalibrated candidate scores.",
                    ),
                    (
                        "jev-latest",
                        "Compatibility alias for "
                        + settings.model_name
                        + "; NOT the official Jev model.",
                    ),
                ]
            ]
        }

    @app.post("/v1/systemone")
    async def systemone(request: Request):
        authorize(request)
        if not app.state.ready:
            raise ServingError(503, "not_ready", "Backends are not ready")
        if app.state.inflight >= settings.max_inflight:
            raise ServingError(429, "busy", "Gateway in-flight budget exceeded")
        app.state.inflight += 1
        try:
            async with asyncio.timeout(settings.timeout):
                raw = bytearray()
                async for chunk in request.stream():
                    raw.extend(chunk)
                    if len(raw) > settings.max_body_bytes:
                        raise ServingError(
                            413, "request_too_large", "Request byte budget exceeded"
                        )
                try:
                    body = json.loads(
                        raw,
                        object_pairs_hook=unique_object,
                        parse_constant=reject_constant,
                    )
                    if not isinstance(body, dict):
                        raise ValueError
                except (ValueError, TypeError, RecursionError, UnicodeError):
                    raise ServingError(
                        422, "invalid_json", "Malformed JSON object"
                    ) from None
                effort = body.get("effort", "high")
                if effort not in ("low", "high"):
                    raise ServingError(
                        422, "invalid_effort", "effort must be low or high"
                    )
                questions = prepare_request(
                    body,
                    model_name=settings.model_name,
                    max_questions=settings.max_questions,
                    allow_score=True,
                )
                compiled = [
                    compile_question(q, tokenizer, max_tokens=settings.max_tokens)
                    for q in questions
                ]
                url, model, depth = (
                    (settings.low_url, settings.low_model, 16)
                    if effort == "low"
                    else (settings.high_url, settings.high_model, 32)
                )
                answers = {}
                for item in compiled:
                    if not item.prompt_ids:
                        answers[item.question.key] = answer(item.question, [1.0])
                        continue
                    payload = {
                        "model": model,
                        "prompt": item.prompt_ids,
                        "max_tokens": 1,
                        "n": 1,
                        "stream": False,
                        "echo": False,
                        "add_special_tokens": False,
                        "logprobs": len(item.candidate_token_ids),
                        "logprob_token_ids": item.candidate_token_ids,
                        "return_tokens_as_token_ids": True,
                        "temperature": 1.0,
                        "top_p": 1.0,
                        "top_k": -1,
                        "min_p": 0.0,
                        "frequency_penalty": 0.0,
                        "presence_penalty": 0.0,
                        "repetition_penalty": 1.0,
                    }
                    response = await app.state.client.post(
                        endpoint(url, "/v1/completions"), json=payload
                    )
                    if response.status_code != 200:
                        raise ServingError(
                            502,
                            "backend_error",
                            f"Inference backend returned HTTP {response.status_code}",
                        )
                    try:
                        data = json.loads(
                            response.content,
                            object_pairs_hook=unique_object,
                            parse_constant=reject_constant,
                        )
                    except (ValueError, TypeError, UnicodeError, RecursionError):
                        raise ServingError(
                            502, "backend_json", "Backend returned invalid JSON"
                        ) from None
                    answers[item.question.key] = answer(
                        item.question,
                        probabilities(parse_scores(data, item.candidate_token_ids)),
                    )
                result = {
                    "model": settings.model_name,
                    "answers": answers,
                    "usage": {
                        "input_tokens": sum(len(c.prompt_ids) for c in compiled),
                        "output_tokens": sum(bool(c.prompt_ids) for c in compiled),
                        "billing_units": len(compiled),
                    },
                }
                return JSONResponse(
                    result,
                    headers={
                        "x-openjev-effort": effort,
                        "x-openjev-executed-layers": str(
                            depth if any(c.prompt_ids for c in compiled) else 0
                        ),
                        "x-openjev-engine-model": model,
                        "x-typesafe-request-id": str(uuid.uuid4()),
                    },
                )
        except (TimeoutError, httpx.TimeoutException):
            raise ServingError(
                504, "deadline_exceeded", "Inference request exceeded timeout"
            ) from None
        except httpx.HTTPError:
            raise ServingError(
                502, "backend_transport", "Inference backend connection failed"
            ) from None
        finally:
            app.state.inflight -= 1

    return app


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=30013)
    args = parser.parse_args()
    from transformers import AutoTokenizer
    import uvicorn

    settings = Settings.from_env()
    if (
        Path(os.environ["JEV_MODEL_PATH"]).resolve()
        != Path(settings.high_model_path).resolve()
    ):
        raise ValueError("Tokenizer must come from the configured full-depth model")
    tokenizer = AutoTokenizer.from_pretrained(
        os.environ["JEV_MODEL_PATH"], local_files_only=True
    )
    uvicorn.run(create_app(settings, tokenizer), host=args.host, port=args.port)


if __name__ == "__main__":
    main()