import os
import time
import torch
from fastapi import FastAPI, HTTPException
from fastapi.responses import FileResponse
from pydantic import BaseModel
from transformers import AutoModelForSequenceClassification, AutoTokenizer

REPO = "AlexWortega/openjev"
MODELS = {
    "0.8b": "qwen3.5-0.8b-nli-v2s-long",
}

app = FastAPI(title="openjev-bench")
_loaded = {}


@app.get("/")
def index():
    return FileResponse(os.path.join(os.path.dirname(__file__), "public", "index.html"))


def load_all():
    for key, sub in MODELS.items():
        tok = AutoTokenizer.from_pretrained(REPO, subfolder=sub)
        # fp32 on CPU is much faster than bf16 (no AMX -> bf16 is emulated);
        # 4B fp32 = 16GB would OOM on cpu-basic (16GB), so keep bf16 for 4b only.
        dtype = torch.float32 if key == "0.8b" else torch.bfloat16
        model = AutoModelForSequenceClassification.from_pretrained(REPO, subfolder=sub, dtype=dtype)
        tok.padding_side = "right"
        model.config.get_text_config().pad_token_id = tok.pad_token_id
        model.eval()
        _loaded[key] = (model, tok)
        print(f"[openjev] loaded {key} ({sub}, dtype={dtype})", flush=True)


@app.on_event("startup")
def _startup():
    load_all()


class Pair(BaseModel):
    premise: str
    hypothesis: str


class PredictReq(BaseModel):
    pairs: list[Pair]
    model: str = "0.8b"


@app.post("/api/predict")
def predict(req: PredictReq):
    requested = req.model
    if requested not in _loaded:
        if requested == "4b":
            requested = "0.8b"
        else:
            raise HTTPException(422, "model must be one of: 0.8b, 4b")
    model, tok = _loaded[requested]
    texts = [model.config.nli_template.format(premise=p.premise, hypothesis=p.hypothesis) for p in req.pairs]
    out = []
    BATCH = 8
    t0 = time.time()
    with torch.no_grad():
        for i in range(0, len(texts), BATCH):
            enc = tok(texts[i:i + BATCH], padding=True, return_tensors="pt")
            logits = model(**enc).logits.float().softmax(-1)
            for probs in logits.tolist():
                out.append({"contradiction": probs[0], "entailment": probs[1], "neutral": probs[2]})
    dt = time.time() - t0
    return {
        "probs": out,
        "elapsed_s": round(dt, 2),
        "ms_per_pair": round(dt * 1000 / max(1, len(out)), 1),
        "model": requested,
        "requestedModel": requested if requested == req.model else req.model,
    }


@app.get("/api/eval")
def eval_mnli(model: str = "0.8b", n: int = 200):
    if model not in MODELS:
        raise HTTPException(422, "model must be one of: " + ", ".join(MODELS))
    model_obj, tok = _loaded[model]
    from datasets import load_dataset

    ds = load_dataset("nyu-mll/glue", "mnli", split="validation_matched")
    ds = ds.shuffle(seed=42).select(range(min(n, len(ds))))

    glue_to_openjev = {0: 1, 1: 2, 2: 0}
    texts = [model_obj.config.nli_template.format(premise=row["premise"], hypothesis=row["hypothesis"]) for row in ds]
    golds = [glue_to_openjev[row["label"]] for row in ds]

    preds = []
    BATCH = 8
    t0 = time.time()
    with torch.no_grad():
        for i in range(0, len(texts), BATCH):
            enc = tok(texts[i:i + BATCH], padding=True, return_tensors="pt")
            logits = model_obj(**enc).logits
            preds.extend(logits.argmax(-1).tolist())
    dt = time.time() - t0

    correct = sum(1 for p, g in zip(preds, golds) if p == g)
    published = {"0.8b": 0.869, "4b": 0.904}
    return {
        "model": model,
        "n": len(texts),
        "accuracy": round(correct / len(texts), 4),
        "published_mnli_m": published.get(model),
        "elapsed_s": round(dt, 2),
        "ms_per_pair": round(dt * 1000 / len(texts), 1),
    }


@app.get("/api/health")
def health():
    return {"status": "ok", "models": {k: {"loaded": k in _loaded, "subfolder": v} for k, v in MODELS.items()}}


@app.get("/api/sysinfo")
def sysinfo():
    def _read(path):
        try:
            with open(path) as f:
                return f.read().strip()
        except OSError:
            return None

    cpu_model = None
    raw = _read("/proc/cpuinfo")
    if raw:
        for line in raw.splitlines():
            if line.lower().startswith("model name"):
                cpu_model = line.split(":", 1)[1].strip()
                break

    mem = {}
    memraw = _read("/proc/meminfo")
    if memraw:
        for line in memraw.splitlines():
            if line.startswith(("MemTotal", "MemAvailable", "MemFree")):
                k, v = line.split(":", 1)
                mem[k] = v.strip()

    return {
        "cpu_model": cpu_model,
        "cpu_affinity": len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else os.cpu_count(),
        "cgroup_cpu_max": _read("/sys/fs/cgroup/cpu.max"),
        "cgroup_memory_max": _read("/sys/fs/cgroup/memory.max"),
        "meminfo": mem,
        "loadavg": _read("/proc/loadavg"),
        "self_vm_rss": next((l.split(":", 1)[1].strip() for l in (_read("/proc/self/status") or "").splitlines() if l.startswith("VmRSS")), None),
    }