import hashlib
import time
"""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
import time
from collections import defaultdict, deque

from fastapi import FastAPI, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware

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 = os.getenv("DECIDE_NAME", "darwin-35b-ztc-vision")
Q: "queue.Queue[tuple[dict, Future]]" = queue.Queue()


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


@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}


API_KEY = os.getenv("DECIDE_API_KEY")
# Demo key for a browser-only (static) page: the key is visible to viewers, so it is
# limited to listed origins, a per-client and a global rate, and small requests.
DEMO_KEY = os.getenv("DECIDE_DEMO_KEY")
DEMO_ORIGINS = [o for o in os.getenv("DECIDE_DEMO_ORIGINS", "").split(",") if o]
DEMO_PER_MIN = int(os.getenv("DECIDE_DEMO_PER_MIN", "40"))
DEMO_GLOBAL_PER_MIN = int(os.getenv("DECIDE_DEMO_GLOBAL_PER_MIN", "240"))
HITS = defaultdict(deque)
if DEMO_ORIGINS:
    app.add_middleware(CORSMiddleware, allow_origins=DEMO_ORIGINS, allow_methods=["POST", "GET", "OPTIONS"],
                       allow_headers=["Authorization", "Content-Type"], max_age=600)


def demo_check(request, body):
    if DEMO_ORIGINS and request.headers.get("origin", "") not in DEMO_ORIGINS:
        raise HTTPException(403, "This key only works from the demo page.")
    now = time.time()
    client = request.headers.get("cf-connecting-ip") or (request.client.host if request.client else "?")
    for name, limit in ((client, DEMO_PER_MIN), ("*", DEMO_GLOBAL_PER_MIN)):
        q = HITS[name]
        while q and now - q[0] > 60:
            q.popleft()
        if len(q) >= limit:
            raise HTTPException(429, "Too many requests. Wait a minute and try again.")
    for name in (client, "*"):
        HITS[name].append(now)
    images = body.get("images") or []
    if len(images) > 4 or any(len(i) > 4_000_000 for i in images):
        raise HTTPException(413, "Up to 4 images, each under about 3 MB.")
    if len(body.get("questions") or {}) > 4:
        raise HTTPException(413, "Up to 4 questions per request.")


@app.post("/v1/systemone")
async def system_one(request: Request):
    auth = request.headers.get("authorization", "")
    demo = bool(DEMO_KEY) and auth == "Bearer " + DEMO_KEY
    if API_KEY and auth != "Bearer " + API_KEY and not demo:
        raise HTTPException(401, "Invalid API key.")
    body = await request.json()
    if demo:
        demo_check(request, body)
    questions = body.get("questions") or {}
    images = body.get("images") or []
    if images and model.text_only:
        raise HTTPException(422, "This model is text-only; images are not supported.")
    if len(images) > 4:
        raise HTTPException(422, "At most 4 images per request.")
    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()
        row = {"state": body.get("state"), "question": question}
        if images:
            row["images"] = images
        Q.put((row, 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))
        bad = not all(math.isfinite(v) and v >= 0 for v in values) or sum(values) <= 0
        flat = (not bad) and len(values) > 1 and max(values) - min(values) < 1e-9
        if bad or flat:
            key_hash = hashlib.sha256(json.dumps([body.get('state'), questions[key], len(images)], sort_keys=True, default=str).encode()).hexdigest()[:16]
            with open(os.getenv('INVALID_LOG', 'invalid_outputs.jsonl'), 'a') as fh:
                fh.write(json.dumps({'time': time.time(), 'request': key_hash, 'question': key, 'kind': 'invalid' if bad else 'uniform', 'values': [repr(v) for v in values]}) + '\n')
        if bad:
            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")