"""Masked-diffusion chat — server-side CPU inference.

The model (LFM2.5-Encoder-350M-Diffusion, f32 safetensors) is loaded once and runs
on the server's CPU. Nothing is downloaded to the visitor's browser: the page streams
denoising frames from /api/generate as newline-delimited JSON and animates them,
exactly like the old in-browser WebGPU build did.
"""

import json
import os
import threading

import torch
from flask import Flask, Response, request, send_from_directory, stream_with_context
from tokenizers import Tokenizer

import lfm2_diffusion as L

APP_DIR = os.path.dirname(os.path.abspath(__file__))
# Match threads to the cores actually granted to the container (cgroup quota),
# not the host's core count — oversubscription badly slows CPU matmuls.
_CPUS = L.effective_cpus()
torch.set_num_threads(_CPUS)
try:
    torch.set_num_interop_threads(1)
except RuntimeError:
    pass

app = Flask(__name__)

# ---- load model + tokenizer once, at import (per gunicorn worker) -----------
HF_TOKEN = os.environ.get("HF_TOKEN")
print(f"[boot] pulling {L.MODEL_ID} from the Hub, loading on {_CPUS} CPU thread(s)…", flush=True)
from huggingface_hub import hf_hub_download
TOK = Tokenizer.from_file(hf_hub_download(L.MODEL_ID, "tokenizer.json", token=HF_TOKEN))
MODEL = L.Lfm2Diffusion.from_hub(token=HF_TOKEN)
print("[boot] calibrating + quantizing (smoothed int8)…", flush=True)
MODEL.quantize_int8([TOK.encode(L.render_chat_prompt(p)).ids for p in L.CALIB_PROMPTS])
_GEN_LOCK = threading.Lock()
# warm up the graph so the first real request is not the slow one
list(L.diffuse(MODEL, TOK.encode(L.render_chat_prompt("Hi")).ids,
               max_new=8, steps=4, block_size=8))
print("[boot] model ready", flush=True)


def _clampi(v, lo, hi, default):
    try:
        return max(lo, min(hi, int(float(v))))
    except (TypeError, ValueError):
        return default


@app.route("/")
def index():
    return send_from_directory(APP_DIR, "index.html")


@app.route("/style.css")
def style():
    return send_from_directory(APP_DIR, "style.css")


@app.route("/lliquid.gif")
def logo():
    return send_from_directory(APP_DIR, "lliquid.gif")


@app.route("/api/ready")
def ready():
    return {"ready": MODEL is not None}


@app.route("/api/generate")
def generate():
    prompt = (request.args.get("prompt") or "").strip()
    if not prompt:
        return {"error": "empty prompt"}, 400
    opts = dict(
        max_new=_clampi(request.args.get("mnt"), 16, 192, 64),
        steps=_clampi(request.args.get("steps"), 8, 128, 32),
        block_size=_clampi(request.args.get("blk"), 8, 64, 16),
        temperature=_clampi(request.args.get("temp"), 0, 100, 0) / 100.0,
        tau=_clampi(request.args.get("tau"), 50, 100, 90) / 100.0,
    )
    prompt_ids = TOK.encode(L.render_chat_prompt(prompt)).ids

    @stream_with_context
    def frames():
        # Serialize generation: the CPU model handles one request at a time.
        if not _GEN_LOCK.acquire(blocking=False):
            yield json.dumps({"error": "The model is busy with another request. Try again in a moment."}) + "\n"
            return
        try:
            last = None
            for fr in L.diffuse(MODEL, prompt_ids, **opts):
                last = fr["gen_ids"]
                tokens = [None if t == L.MASK_ID
                          else TOK.decode([t], skip_special_tokens=True)
                          for t in fr["gen_ids"]]
                yield json.dumps({
                    "step": fr["step"],
                    "tokens": tokens,
                    "revealed": sorted(fr["revealed"]),
                }) + "\n"
            gen = [t for t in (last or []) if t != L.MASK_ID]
            text = L.trim(TOK.decode(gen, skip_special_tokens=True))
            yield json.dumps({"done": True, "text": text or "(empty)"}) + "\n"
        except Exception as exc:  # surface errors to the UI
            yield json.dumps({"error": str(exc)}) + "\n"
        finally:
            _GEN_LOCK.release()

    return Response(frames(), mimetype="application/x-ndjson", headers={
        "Cache-Control": "no-cache, no-transform",
        "X-Accel-Buffering": "no",
    })


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=int(os.environ.get("PORT", 7860)), threaded=True)