"""POST /v1/systemone server for FINAL-Bench/Darwin-27B-ZTC (System One request shape).

    ZTC_MODEL=<snapshot path or FINAL-Bench/Darwin-27B-ZTC> PORT=8000 python ztc_server.py

Every question is answered with one forward pass of the backbone and the decision readout
(no generated tokens). Questions from concurrent requests are batched together on the GPU
(at most ZTC_BATCH questions and ZTC_TOKEN_BUDGET estimated tokens per batch). A question
longer than ZTC_MAX_LENGTH tokens is refused with HTTP 422 (never truncated).
"""

import asyncio
import json
import math
import os
import queue
import sys
import threading
from concurrent.futures import Future
from pathlib import Path

import torch
import uvicorn
from fastapi import FastAPI, HTTPException, Request

MODEL = os.getenv("ZTC_MODEL", "FINAL-Bench/Darwin-27B-ZTC")
MAX_LENGTH = int(os.getenv("ZTC_MAX_LENGTH", "32768"))
BATCH = int(os.getenv("ZTC_BATCH", "8"))
TOKEN_BUDGET = int(os.getenv("ZTC_TOKEN_BUDGET", "16000"))

path = Path(MODEL)
if not path.exists():
    from huggingface_hub import snapshot_download
    path = Path(snapshot_download(MODEL, revision=os.getenv("ZTC_REVISION")))
sys.path.insert(0, str(path))
from autojev.model import DecisionModel, answer  # noqa: E402

model = DecisionModel(checkpoint=path).eval()
NAME = "darwin-27b-ztc"
Q: "queue.Queue[tuple[dict, Future]]" = queue.Queue()


def estimate(row):
    return len(json.dumps([row["state"], row["question"]], ensure_ascii=False)) // 3 + 192


@torch.inference_mode()
def probabilities(rows):
    batch = model.prepare(rows, max_length=MAX_LENGTH)
    values = (model(batch) / model.temperature).softmax(-1).cpu().tolist()
    return [v[:c] for v, c in zip(values, batch.counts)]


def run_single(row, future):
    try:
        future.set_result(probabilities([row])[0])
    except ValueError as error:
        future.set_exception(ValueError(str(error)))
    except Exception as error:  # out of memory or other runtime failure
        torch.cuda.empty_cache()
        future.set_exception(RuntimeError(str(error)[:300]))


def worker():
    while True:
        items = [Q.get()]
        while len(items) < BATCH * 4:
            try:
                items.append(Q.get(timeout=0.005))
            except queue.Empty:
                break
        items.sort(key=lambda item: estimate(item[0]))
        i = 0
        while i < len(items):
            group = [items[i]]
            tokens = estimate(items[i][0])
            while i + len(group) < len(items) and len(group) < BATCH and tokens + estimate(items[i + len(group)][0]) <= TOKEN_BUDGET:
                tokens += estimate(items[i + len(group)][0])
                group.append(items[i + len(group)])
            i += len(group)
            try:
                for (row, future), values in zip(group, probabilities([g[0] for g in group])):
                    future.set_result(values)
            except Exception:
                torch.cuda.empty_cache()
                for row, future in group:
                    if not future.done():
                        run_single(row, future)


threading.Thread(target=worker, daemon=True).start()
app = FastAPI(title="Darwin-27B-ZTC System One server")


@app.get("/health")
def health():
    return {"status": "ready", "model": NAME, "temperature": model.temperature}


@app.post("/v1/systemone")
async def system_one(request: Request):
    body = await request.json()
    questions = body.get("questions") or {}
    if body.get("images"):
        raise HTTPException(422, "This model is text-only; images are not supported.")
    futures = {}
    for key, question in questions.items():
        if question.get("type") not in ("choice", "noul", "score"):
            raise HTTPException(422, f"Unsupported question type: {question.get('type')}")
        future = Future()
        Q.put(({"state": body.get("state"), "question": question}, future))
        futures[key] = future
    answers = {}
    for key, future in futures.items():
        try:
            values = await asyncio.wrap_future(future)
        except ValueError as error:
            raise HTTPException(422, str(error))
        except RuntimeError as error:
            raise HTTPException(500, str(error))
        if not all(math.isfinite(v) and v >= 0 for v in values) or sum(values) <= 0:
            raise HTTPException(500, 'Model produced an invalid probability vector; refusing to substitute a prediction.')
        answers[key] = answer(questions[key], values)
    return {"model": NAME, "answers": answers}


if __name__ == "__main__":
    uvicorn.run(app, host=os.getenv("HOST", "127.0.0.1"), port=int(os.getenv("PORT", "8000")), log_level="warning")