import os, json, threading, time

# ─── set thread env BEFORE torch import ───────────────────────────
_nt = str(os.cpu_count() or 4)
os.environ.setdefault("OMP_NUM_THREADS", _nt)
os.environ.setdefault("MKL_NUM_THREADS", _nt)

import torch
from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse, JSONResponse, FileResponse
from fastapi.staticfiles import StaticFiles
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer

# ─── CPU thread opts ────────────────────────────────────────────────
torch.set_num_threads(os.cpu_count() or 4)
try:
    torch.set_num_interop_threads(max(1, (os.cpu_count() or 4) // 2))
except RuntimeError:
    pass  # can only set once per process

log = lambda m: print(m, flush=True)

# ─── Globals ─────────────────────────────────────────────────────────
MODEL_ID = "Smilyai-labs/Nova-1-Standard-1.3B-Preview"
model = None
tokenizer = None
model_status = "loading"
model_error = ""
load_stage = "starting"

# ─── Background loader ──────────────────────────────────────────────
def load_model():
    global model, tokenizer, model_status, model_error, load_stage
    try:
        load_stage = "tokenizer"
        log("[nova] loading tokenizer …")
        tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
        log("[nova] tokenizer ready ✓")

        load_stage = "weights"
        log("[nova] loading model weights (float32, no quant) …")
        model = AutoModelForCausalLM.from_pretrained(
            MODEL_ID,
            trust_remote_code=True,
            dtype=torch.float32,        # float32 — safest/fastest on CPU
            low_cpu_mem_usage=True,
        )
        n_params = sum(p.numel() for p in model.parameters()) / 1e9
        log(f"[nova] weights loaded ✓ ({n_params:.2f}B params)")

        # ── NO int8 quant ──
        # For batch_size=1 on CPU the per-op quantize/dequantize overhead
        # makes dynamic int8 SLOWER than plain float32.
        # The bottleneck is compute (use_cache=False) not memory bandwidth,
        # so quant doesn't help — it just adds latency.
        log("[nova] skipping int8 quant (slower for bs=1 CPU inference)")

        model.eval()
        model_status = "ready"
        load_stage = "ready"
        log("[nova] ✅ MODEL READY — go ahead and chat!")
    except Exception as e:
        model_status = "error"
        model_error = str(e)
        load_stage = "error"
        log(f"[nova] ❌ FATAL: {e}")

threading.Thread(target=load_model, daemon=True).start()

# ─── FastAPI app ────────────────────────────────────────────────────
app = FastAPI()
app.mount("/static", StaticFiles(directory="static"), name="static")

@app.get("/")
async def root():
    return FileResponse("static/index.html")

@app.get("/api/health")
async def health():
    return {"status": model_status, "stage": load_stage, "error": model_error}

# ─── Chat endpoint (SSE streaming) ──────────────────────────────────
@app.post("/api/chat")
async def chat(request: Request):
    if model_status != "ready":
        return JSONResponse(
            {"error": f"Model not ready (stage={load_stage})"},
            status_code=503,
        )

    body = await request.json()
    # ── BRUTAL context truncation ──
    # Nova-1 has use_cache=False baked in → every token recomputes the
    # ENTIRE context. Shorter context = exponentially faster.
    # Keep only last 2 user turns.
    messages = body.get("messages", [])[-2:]
    temperature = float(body.get("temperature", 0.7))
    max_new_tokens = int(body.get("max_new_tokens", 128))
    top_p = float(body.get("top_p", 0.9))

    chat_messages = [{"role": "system",
                      "content": "You are Nova, a helpful, honest AI assistant. Keep answers concise."}]
    chat_messages.extend(messages)

    input_ids = tokenizer.apply_chat_template(
        chat_messages,
        add_generation_prompt=True,
        return_tensors="pt",
    )

    streamer = TextIteratorStreamer(
        tokenizer, skip_prompt=True, skip_special_tokens=True,
    )

    do_sample = temperature > 0
    gen_kwargs = dict(
        input_ids=input_ids,       # ← FIXED: was inputs= (caused AttributeError)
        max_new_tokens=max_new_tokens,
        streamer=streamer,
        use_cache=False,           # nova1 architecture — no KV cache possible
        do_sample=do_sample,
        num_beams=1,               # greedy/sample only — no beam search
        pad_token_id=tokenizer.eos_token_id,
    )
    if do_sample:
        gen_kwargs["temperature"] = max(temperature, 0.01)
        gen_kwargs["top_p"] = top_p

    @torch.inference_mode()
    def _generate():
        model.generate(**gen_kwargs)

    thread = threading.Thread(target=_generate, daemon=True)
    thread.start()

    def event_stream():
        try:
            for token in streamer:
                if token:
                    yield f"data: {json.dumps({'token': token})}\n\n"
        except Exception as e:
            yield f"data: {json.dumps({'error': str(e)})}\n\n"
        yield "data: [DONE]\n\n"

    return StreamingResponse(
        event_stream(),
        media_type="text/event-stream",
        headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
    )