"""Nova Agentic Brain · Qwen3-Coder-30B-A3B-Instruct on HF ZeroGPU.

Backs `/nova/agentic` on the nova-brain Cloudflare Worker. Streaming OpenAI-
compatible chat for tool-call / code / dispatch turns of the Sovereign CEO
Assistant.
"""
from __future__ import annotations

import json
import time
from threading import Thread
from typing import Iterable

import gradio as gr
import spaces
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer

MODEL_ID = "Qwen/Qwen3-Coder-30B-A3B-Instruct"

print(f"loading {MODEL_ID} (CPU init · GPU on call)...", flush=True)
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
    device_map="auto",
    trust_remote_code=True,
)
print("model ready · waiting for ZeroGPU dispatch", flush=True)


def _normalize_messages(raw) -> list[dict[str, str]]:
    if isinstance(raw, str):
        try:
            raw = json.loads(raw)
        except Exception:
            raw = [{"role": "user", "content": raw}]
    out = []
    for m in raw:
        if isinstance(m, dict) and "role" in m and "content" in m:
            out.append({"role": str(m["role"]), "content": str(m["content"])})
    return out or [{"role": "user", "content": ""}]


@spaces.GPU(duration=120)
def chat(messages_json: str, max_tokens: int = 512, temperature: float = 0.6) -> Iterable[str]:
    """Streaming chat. messages_json = JSON array of {role, content}."""
    msgs = _normalize_messages(messages_json)
    prompt = tokenizer.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)

    gen_kwargs = dict(
        **inputs,
        max_new_tokens=int(max_tokens),
        temperature=float(temperature),
        do_sample=temperature > 0.01,
        streamer=streamer,
        pad_token_id=tokenizer.eos_token_id,
    )
    t = Thread(target=model.generate, kwargs=gen_kwargs)
    t.start()

    accumulated = ""
    started = time.time()
    for chunk in streamer:
        accumulated += chunk
        yield accumulated
    t.join()
    print(f"chat done · {len(accumulated)}ch · {time.time()-started:.1f}s", flush=True)


@spaces.GPU(duration=60)
def chat_oneshot(messages_json: str, max_tokens: int = 256, temperature: float = 0.3) -> str:
    """Non-streaming version for fast tool-call decisions."""
    msgs = _normalize_messages(messages_json)
    prompt = tokenizer.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    out = model.generate(
        **inputs,
        max_new_tokens=int(max_tokens),
        temperature=float(temperature),
        do_sample=temperature > 0.01,
        pad_token_id=tokenizer.eos_token_id,
    )
    text = tokenizer.decode(out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
    return text


with gr.Blocks(title="Nova Qwen3-Coder Agentic Brain") as demo:
    gr.Markdown("# Nova Qwen3-Coder · Agentic Brain")
    gr.Markdown("Backs `/nova/agentic` on the nova-brain Cloudflare Worker.")

    with gr.Tab("Streaming chat"):
        msgs_in = gr.Textbox(label="messages_json", lines=4, value='[{"role":"user","content":"Hello"}]')
        max_t = gr.Slider(32, 2048, value=512, step=32, label="max_tokens")
        temp = gr.Slider(0.0, 1.5, value=0.6, step=0.05, label="temperature")
        out = gr.Textbox(label="output (streaming)", lines=10)
        gr.Button("Run").click(chat, [msgs_in, max_t, temp], out)

    with gr.Tab("One-shot"):
        msgs_in2 = gr.Textbox(label="messages_json", lines=4, value='[{"role":"user","content":"Reply: OK"}]')
        max_t2 = gr.Slider(8, 1024, value=256, step=16, label="max_tokens")
        temp2 = gr.Slider(0.0, 1.5, value=0.3, step=0.05, label="temperature")
        out2 = gr.Textbox(label="output", lines=6)
        gr.Button("Run").click(chat_oneshot, [msgs_in2, max_t2, temp2], out2)

if __name__ == "__main__":
    demo.queue(max_size=10).launch(server_name="0.0.0.0", server_port=7860)