"""Jev-shaped request/response wrapper around HevForDecisions, plus an optional FastAPI server."""
from __future__ import annotations

import argparse
import json
import threading

from .schema import Question

MAX_QUESTIONS = 32
# A question's canvas entry (its id, the JSON around it and the answer slot) takes 6 or more DiffusionGemma tokens,
# so 4 is a safe floor for the early check below.
MIN_CANVAS_TOKENS_PER_QUESTION = 4


def max_questions(model) -> int:
    """Upper bound on questions per request for `model`'s answer canvas (32 when the model has no canvas_length).

    Every question needs more than MIN_CANVAS_TOKENS_PER_QUESTION canvas tokens, so this never rejects a request that
    would fit. With short ids a 64-token canvas holds about 8 questions (9 one-letter ids, 7 word ids); the exact
    check happens when the canvas is built and raises ValueError, which the server turns into a 422.
    """
    canvas = getattr(getattr(model, "config", None), "canvas_length", None)
    if not isinstance(canvas, int) or isinstance(canvas, bool) or canvas <= 0:
        return MAX_QUESTIONS
    return max(1, min(MAX_QUESTIONS, canvas // MIN_CANVAS_TOKENS_PER_QUESTION))


def validate_request(request: dict, max_q: int = MAX_QUESTIONS) -> None:
    if not isinstance(request, dict) or "state" not in request:
        raise ValueError("request needs a 'state' (string, object or array)")
    qs = request.get("questions")
    if not isinstance(qs, dict) or not qs:
        raise ValueError("request needs a non-empty 'questions' mapping")
    for qid, spec in qs.items():
        if not isinstance(spec, dict):
            raise ValueError(f"question {qid!r} must be an object")
        try:
            Question.from_jev(qid, spec)
        except ValueError as err:
            raise ValueError(f"question {qid!r}: {err}") from err
    n = len(qs)
    if n > max_q:
        raise ValueError(f"at most {max_q} questions per request fit this model's answer canvas, got {n}")


def _state_text(state) -> str:
    return state if isinstance(state, str) else json.dumps(state, ensure_ascii=False)


def decide(model, request: dict) -> dict:
    validate_request(request, max_questions(model))
    resp = model.decide(_state_text(request["state"]), request["questions"])
    resp["model"] = request.get("model", "hev")
    return resp


def build_app(model):
    """FastAPI app serving `model` at POST /v1/systemone and GET /health.

    FastAPI runs plain `def` endpoints in a thread pool, and HevForDecisions.decide() is not
    thread-safe (it shares one hook-capture buffer on the backbone), so calls to the model are
    serialised with a lock: concurrent requests queue and are answered one at a time.
    """
    from fastapi import FastAPI, HTTPException

    app = FastAPI(title="Hev", version="0.1.0")
    model_lock = threading.Lock()

    @app.get("/health")
    def health():
        return {"ok": True}

    @app.post("/v1/systemone")
    def systemone(request: dict):
        try:
            with model_lock:
                return decide(model, request)
        except ValueError as err:
            raise HTTPException(status_code=422, detail=str(err)) from err

    return app


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--repo", required=True)
    ap.add_argument("--port", type=int, default=8080)
    ap.add_argument("--host", default="127.0.0.1")
    ap.add_argument("--load-in-4bit", action="store_true")
    ap.add_argument("--request", default=None, help="JSON file: answer once and exit")
    args = ap.parse_args()
    from .modeling_hev import HevForDecisions
    model = HevForDecisions.from_pretrained(args.repo, load_in_4bit=args.load_in_4bit)
    if args.request:
        with open(args.request, "r", encoding="utf-8") as f:
            print(json.dumps(decide(model, json.load(f)), indent=2))
        return
    import uvicorn
    uvicorn.run(build_app(model), host=args.host, port=args.port)


if __name__ == "__main__":
    main()