"""Inference server for node12 behind ``api.askdol.com`` at ``/naver`` (classifier-plan §6.2 option (a), §6.4).

Run as ``python -m training.serve --package DIR --base DIR [--host 127.0.0.1 --port 8012]``. ``--package`` is a
pinned snapshot of the private HF repo (``export_hf`` layout: adapter/, heads.safetensors, taxonomy.json,
temperatures.json, config.json, MODEL_REVISION); ``--base`` is the local Gemma 4 12B checkpoint it was trained on.

Endpoints (Bearer key from ``NAVER_PSY_CLASSIFIER_API_KEY``, compared in constant time, never logged):
- ``GET /naver/health`` -> repo, revision, model_revision, taxonomy (canonical), temperatures, fields, max_batch;
- ``POST /naver/predict`` ``{"inputs": [{"id", "segments": [{"kind", "text"}]}]}`` -> per input raw logits,
  calibrated gated probabilities and labels, ``n_tokens`` and ``truncated``. Post text is never echoed or logged.

The torch-free core ``Service`` takes a ``predict_fn(encoded ids) -> list of {field: logits}``, so it is tested
without a model; ``load_service`` builds the real one.
"""
from __future__ import annotations

import argparse
import hmac
import json
import logging
import os
import re
from pathlib import Path
from typing import Any, Callable, Mapping

from .calibrate import outputs
from .formats import encode
from .taxonomy import Taxonomy

log = logging.getLogger("training.serve")
API_KEY_ENV = "NAVER_PSY_CLASSIFIER_API_KEY"
REVISION_ENV = "NAVER_PSY_MODEL_REVISION"
RENDER_VERSION = "nsb-input-v1"
PREFIX = "/naver"
REVISION_RE = re.compile(r"^[^@\s]+@[0-9a-f]{40}$")
MAX_SEGMENTS = 16
MAX_TEXT = 20_000

PredictFn = Callable[[list[list[int]]], list[dict[str, list[float]]]]


class RequestError(ValueError):
    """A malformed request (HTTP 400/413); the message never contains post text."""

    def __init__(self, message: str, status: int = 400) -> None:
        super().__init__(message)
        self.status = status


class Service:
    """Encode, predict, calibrate and gate one batch (torch-free)."""

    def __init__(self, predict_fn: PredictFn, tokenize: Callable[[str], list[int]], bos: int | None, tax: Taxonomy,
                 temps: Mapping[str, float], model_revision: str, repo: str | None = None,
                 revision: str | None = None, max_len: int = 2048, max_batch: int = 64,
                 render_version: str = RENDER_VERSION) -> None:
        if not REVISION_RE.match(model_revision):
            raise ValueError("model_revision must be '<repo>@<40-hex commit>'")
        r, c = model_revision.split("@")
        if (repo and repo != r) or (revision and revision != c):
            raise ValueError("repo/revision differ from model_revision")
        if set(temps) != set(tax.names):
            raise ValueError("temperatures do not cover the taxonomy fields")
        self.predict_fn, self.tokenize, self.bos, self.tax = predict_fn, tokenize, bos, tax
        self.temps = {f: float(temps[f]) for f in tax.names}
        self.model_revision, self.repo, self.revision = model_revision, r, c
        self.max_len, self.max_batch, self.render_version = max_len, max_batch, render_version

    def health(self) -> dict[str, Any]:
        return {"status": "ok", "repo": self.repo, "revision": self.revision, "model_revision": self.model_revision,
                "taxonomy_version": self.tax.version, "render_version": self.render_version,
                "taxonomy": self.tax.to_json(), "temperatures": self.temps, "fields": self.tax.names,
                "max_batch": self.max_batch, "max_len": self.max_len}

    def _validate(self, inputs: Any) -> list[tuple[str, list[dict]]]:
        if not isinstance(inputs, list) or not inputs:
            raise RequestError("inputs must be a non-empty list")
        if len(inputs) > self.max_batch:
            raise RequestError(f"batch of {len(inputs)} exceeds max_batch {self.max_batch}", 413)
        out, seen = [], set()
        for i, item in enumerate(inputs):
            if not isinstance(item, Mapping) or not isinstance(item.get("id"), (str, int)):
                raise RequestError(f"input {i}: needs an id")
            rid = str(item["id"])
            if rid in seen:
                raise RequestError(f"input {i}: duplicate id")
            seen.add(rid)
            segs = item.get("segments")
            if not isinstance(segs, list) or not segs or len(segs) > MAX_SEGMENTS:
                raise RequestError(f"input {i}: segments must be a list of 1-{MAX_SEGMENTS}")
            clean = []
            for s in segs:
                if (not isinstance(s, Mapping) or not isinstance(s.get("kind"), str)
                        or not isinstance(s.get("text"), str) or len(s["text"]) > MAX_TEXT):
                    raise RequestError(f"input {i}: each segment needs a string kind and text (<= {MAX_TEXT} chars)")
                clean.append({"kind": s["kind"], "text": s["text"]})
            out.append((rid, clean))
        return out

    def predict(self, inputs: Any) -> dict[str, Any]:
        items = self._validate(inputs)
        encoded, infos = [], []
        for i, (_, segs) in enumerate(items):
            try:
                ids, info = encode(segs, self.tokenize, self.max_len, self.bos)
            except ValueError:
                raise RequestError(f"input {i}: cannot be encoded within max_len") from None
            encoded.append(ids)
            infos.append(info)
        logits = self.predict_fn(encoded)
        if len(logits) != len(items):
            raise RuntimeError("predict_fn returned a different number of rows")
        results = []
        for (rid, _), lg, info in zip(items, logits, infos):
            row = {f: [float(x) for x in lg[f]] for f in self.tax.names}
            out = outputs(row, self.tax, self.temps)
            results.append({"id": rid, "logits": row, "probs": out["probs"], "labels": out["labels"],
                            "n_tokens": info["n_tokens"], "truncated": bool(info["cut"])})
        return {"model_revision": self.model_revision, "results": results}


def authorised(header: str | None, key: str | None) -> bool:
    """Constant-time Bearer check; no key configured means every request is refused."""
    if not key or not header or not header.startswith("Bearer "):
        return False
    return hmac.compare_digest(header[len("Bearer "):].encode(), key.encode())


def create_app(service: Service, api_key: str | None = None, prefix: str = PREFIX):
    """FastAPI app around a Service (FastAPI is imported here so the core needs nothing)."""
    from fastapi import FastAPI, Request
    from fastapi.responses import JSONResponse

    key = api_key if api_key is not None else os.environ.get(API_KEY_ENV)
    app = FastAPI(title="naver-psy classifier", docs_url=None, redoc_url=None, openapi_url=None)

    def denied() -> JSONResponse:
        return JSONResponse({"error": "unauthorised"}, status_code=401)

    def health(request):
        if not authorised(request.headers.get("authorization"), key):
            return denied()
        return service.health()

    async def predict(request):
        if not authorised(request.headers.get("authorization"), key):
            return denied()
        try:
            body = await request.json()
        except Exception:
            return JSONResponse({"error": "body must be JSON"}, status_code=400)
        try:
            return service.predict(body.get("inputs") if isinstance(body, dict) else None)
        except RequestError as e:
            return JSONResponse({"error": str(e)}, status_code=e.status)
        except Exception as e:  # never echo the input; log the class only
            log.error("predict failed: %s", type(e).__name__)
            return JSONResponse({"error": "internal error"}, status_code=500)

    # Real annotation objects (this module uses postponed annotations, which FastAPI cannot resolve locally).
    health.__annotations__["request"] = Request
    predict.__annotations__["request"] = Request
    app.get(prefix + "/health")(health)
    app.post(prefix + "/predict")(predict)
    return app


def read_revision(package: Path) -> str:
    env = os.environ.get(REVISION_ENV)
    if env:
        return env.strip()
    path = package / "MODEL_REVISION"
    if not path.exists():
        raise SystemExit(f"no model revision: set {REVISION_ENV} or write {path}")
    return path.read_text().strip()


def load_service(package_dir, base_dir, device: str = "cuda", max_batch: int = 64) -> Service:
    """Build the torch model from an HF package snapshot and a local base checkpoint."""
    import torch

    from .data import inference_batches
    from .model import bos_id, build, load_tokenizer, predict, tokenizer_fn

    package = Path(package_dir)
    cfg = json.loads((package / "config.json").read_text())
    cfg["model"] = str(base_dir)
    tax = Taxonomy.load(package / "taxonomy.json")
    temps = json.loads((package / "temperatures.json").read_text())
    dev = torch.device(device)
    model = build(cfg, tax.sizes(), dev, train=False, adapter_dir=package / "adapter",
                  heads_path=package / "heads.safetensors")
    tok = load_tokenizer(str(base_dir))
    max_len = int(cfg.get("max_len", 2048))
    budget = int(cfg.get("serve_token_budget", 65536))

    def predict_fn(encoded: list[list[int]]) -> list[dict[str, list[float]]]:
        batches = inference_batches([len(x) for x in encoded], budget)
        out = predict(model, encoded, tok.pad_token_id, batches, dev, tax.names)
        return [out[i] for i in range(len(encoded))]

    return Service(predict_fn, tokenizer_fn(tok), bos_id(tok), tax, temps, read_revision(package),
                   max_len=max_len, max_batch=max_batch)


def main(argv=None) -> int:
    p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    p.add_argument("--package", required=True, help="Pinned HF package snapshot")
    p.add_argument("--base", required=True, help="Local base checkpoint directory")
    p.add_argument("--host", default="127.0.0.1")
    p.add_argument("--port", type=int, default=8012)
    p.add_argument("--device", default="cuda")
    p.add_argument("--max-batch", type=int, default=64)
    args = p.parse_args(argv)
    if not os.environ.get(API_KEY_ENV):
        raise SystemExit(f"{API_KEY_ENV} is not set; refusing to serve without auth")
    import uvicorn

    logging.basicConfig(level=logging.INFO)
    service = load_service(args.package, args.base, args.device, args.max_batch)
    log.info("serving %s", service.model_revision)
    uvicorn.run(create_app(service), host=args.host, port=args.port, access_log=False)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())