"""Serve a kev decision checkpoint over the `POST /v1/systemone` decision API (ported from autojev).

Scoring goes through kev.decide, the same path as the benchmark engine (kev.bench.KevEngine), so
the endpoint answers exactly as the benchmarked model: same prompt, batching, context limit and
stored temperature. Requests that do not fit are refused with HTTP 413 ("maximum context length"),
never truncated.

  KEV_CHECKPOINT=.../checkpoints/selected KEV_API_KEY=... uv run kev-serve   # env: KEV_HOST, PORT
"""

from __future__ import annotations

import base64
import binascii
import hmac
import os
import threading
import time
import uuid
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from dataclasses import dataclass, field
from datetime import datetime, timezone
from io import BytesIO
from pathlib import Path
from typing import TYPE_CHECKING, Annotated, Literal, cast

from fastapi import Depends, FastAPI, Header, HTTPException, Request
from fastapi.exceptions import RequestValidationError
from fastapi.responses import HTMLResponse, JSONResponse, Response
from PIL import Image, UnidentifiedImageError
from pydantic import BaseModel, ConfigDict, Field, JsonValue, field_validator
from starlette.concurrency import run_in_threadpool
from starlette.middleware.base import RequestResponseEndpoint

from kev.types import Answer, DecisionResponse, JSONValue, Question as DecisionQuestion

if TYPE_CHECKING:
    from kev.model import DecisionModel

type Content = str | dict[str, JsonValue] | list[JsonValue]
MAX_TOKENS = 131072       # same declared context as the benchmark engine
TOKEN_BUDGET = 131072
BATCH_SIZE = 64
QUEUE_SECONDS = 60.0      # a request waits this long for the model before 529


@dataclass
class Service:
    model: DecisionModel | None = None
    name: str = "kev"
    checkpoint: str = ""
    release_date: str = ""
    lock: threading.Lock = field(default_factory=threading.Lock)

    @property
    def aliases(self) -> set[str]:
        return {self.name, "kev-latest"}


service = Service()


class Question(BaseModel):
    model_config = ConfigDict(extra="forbid")
    instructions: Content | None = None


class Choice(Question):
    type: Literal["choice"]
    criteria: dict[str, Content | None] = Field(min_length=1, max_length=255)


class Score(Question):
    type: Literal["score"]
    criteria: list[Content] = Field(min_length=2, max_length=10)


class Noul(Question):
    type: Literal["noul"]
    criteria: dict[Literal["true", "false"], Content | None] | None = None


class EvaluationRequest(BaseModel):
    model_config = ConfigDict(extra="forbid")
    model: str
    state: Content
    questions: dict[str, Annotated[Choice | Score | Noul, Field(discriminator="type")]] = Field(min_length=1)
    images: list[str] = Field(default_factory=list, max_length=4)

    @field_validator("model")
    @classmethod
    def known_model(cls, value: str) -> str:
        if value not in service.aliases:
            raise ValueError(f"Unknown model. Use {service.name} or kev-latest.")
        return value

    @field_validator("images")
    @classmethod
    def valid_images(cls, values: list[str]) -> list[str]:
        for value in values:
            if len(value) > 12_000_000:
                raise ValueError("Each image must be at most 8 MB before base64 encoding.")
            header, separator, encoded = value.partition(",")
            if not separator or header not in {"data:image/png;base64", "data:image/jpeg;base64", "data:image/webp;base64"}:
                raise ValueError("Images must be base64 PNG, JPEG, or WebP data URLs.")
            try:
                content = base64.b64decode(encoded, validate=True)
                if len(content) > 8_000_000:
                    raise ValueError("Each image must be at most 8 MB.")
                with Image.open(BytesIO(content)) as image:
                    if image.width * image.height > 16_000_000:
                        raise ValueError("Each image must have at most 16 million pixels.")
                    if image.format not in {"PNG", "JPEG", "WEBP"}:
                        raise ValueError("Unsupported image format.")
                    image.verify()
            except (binascii.Error, OSError, SyntaxError, UnidentifiedImageError, Image.DecompressionBombError) as error:
                raise ValueError("Invalid image data.") from error
        return values


def authenticate(authorization: str | None = Header(default=None)) -> None:
    key = os.getenv("KEV_API_KEY")
    if key and not hmac.compare_digest((authorization or "").encode(), f"Bearer {key}".encode()):
        raise HTTPException(401, "Missing or invalid API key.", headers={"WWW-Authenticate": "Bearer"})


def model_name(checkpoint: Path) -> str:
    """kev-<run name> for runs/<run>/checkpoints/selected, else kev-<directory name>."""
    resolved = checkpoint.resolve()
    run = resolved.parent.parent.name if resolved.parent.name == "checkpoints" else resolved.name
    return os.getenv("KEV_MODEL_NAME", f"kev-{run}")


@asynccontextmanager
async def lifespan(app: FastAPI) -> AsyncIterator[None]:
    from kev.model import DecisionModel

    checkpoint = os.getenv("KEV_CHECKPOINT")
    if not checkpoint:
        raise RuntimeError("Set KEV_CHECKPOINT to a decision checkpoint directory")
    service.checkpoint = str(Path(checkpoint).resolve())
    service.name = model_name(Path(checkpoint))
    service.model = await run_in_threadpool(DecisionModel, checkpoint=checkpoint, device=os.getenv("KEV_DEVICE") or None)
    modified = (Path(checkpoint) / "decision_config.json").stat().st_mtime
    service.release_date = datetime.fromtimestamp(modified, timezone.utc).date().isoformat()
    try:
        yield
    finally:
        service.model = None


app = FastAPI(title="Kev", version="0.1.0", lifespan=lifespan)


@app.middleware("http")
async def request_metadata(request: Request, call_next: RequestResponseEndpoint) -> Response:
    started, identifier = time.perf_counter(), uuid.uuid4().hex
    response = await call_next(request)
    response.headers["x-request-id"] = identifier
    response.headers["server-timing"] = f"total;dur={(time.perf_counter() - started) * 1000:.1f}"
    return response


@app.exception_handler(RequestValidationError)
async def validation_error(request: Request, error: RequestValidationError) -> JSONResponse:
    return JSONResponse(status_code=422, content={"detail": [
        {"loc": item["loc"], "msg": item["msg"], "type": item["type"]} for item in error.errors()]})


@app.get("/", response_class=HTMLResponse, include_in_schema=False)
def playground() -> str:
    return Path(__file__).with_name("playground.html").read_text()


@app.get("/health", response_model=None)
def health() -> dict[str, JSONValue]:
    model = service.model
    return {"status": "ready" if model is not None else "loading", "model": service.name,
            "checkpoint": service.checkpoint, "temperature": model.temperature if model else None,
            "max_context_tokens": MAX_TOKENS, "authentication": bool(os.getenv("KEV_API_KEY")),
            "modalities": ["text", "image"]}


@app.get("/v1/models", dependencies=[Depends(authenticate)], response_model=None)
def models() -> dict[str, JSONValue]:
    return {"models": [{"name": name, "description": "Kev one-pass typed decisions (text and image).",
                        "release_date": service.release_date} for name in sorted(service.aliases)]}


def predict(model: DecisionModel, body: EvaluationRequest) -> DecisionResponse:
    from kev.decide import decide
    from kev.model import answer

    questions = {key: cast(DecisionQuestion, question.model_dump(exclude_none=True)) for key, question in body.questions.items()}
    distributions, input_tokens = decide(model, body.state, questions, temperature=model.temperature, max_tokens=MAX_TOKENS,
                                         token_budget=TOKEN_BUDGET, batch_size=BATCH_SIZE, images=list(body.images))
    answers: dict[str, Answer] = {key: answer(questions[key], values) for key, values in distributions.items()}
    return {"model": service.name, "answers": answers, "usage": {"input_tokens": input_tokens, "output_tokens": 0}}


@app.post("/v1/systemone", dependencies=[Depends(authenticate)], response_model=None)
async def system_one(body: EvaluationRequest) -> DecisionResponse:
    from kev.decide import CapacityError

    model = service.model
    if model is None:
        raise HTTPException(503, "The model is not ready.")
    if not await run_in_threadpool(service.lock.acquire, True, QUEUE_SECONDS):
        raise HTTPException(529, "The model is busy. Retry shortly.", headers={"Retry-After": "1"})
    try:
        return await run_in_threadpool(predict, model, body)
    except CapacityError as error:
        raise HTTPException(413, str(error)) from error
    except ValueError as error:
        raise HTTPException(422, str(error)) from error
    finally:
        service.lock.release()


def main() -> None:
    import uvicorn

    uvicorn.run("kev.server:app", host=os.getenv("KEV_HOST", "127.0.0.1"), port=int(os.getenv("PORT", "8000")))


if __name__ == "__main__":
    main()