"""AGILLM 4.3 GUI — Hugging Face Space (ZeroGPU).

Runs the single-file runtime in-process: each request calls agillm41.main()
with a patched argv inside the @spaces.GPU window, so the H200 slice is
visible to the loader (--device cuda). The zstd-decompressed checkpoint
cache persists on disk between requests, so only the first request pays
decompression; later requests just load + generate.

Streaming: stdout is captured line-wise and the [STREAM_*] marker protocol
renders NAT as a canvas filling in confidence order.
"""
import json
import os
import queue
import re
import sys
import threading
import time

import gradio as gr
import spaces
from huggingface_hub import hf_hub_download

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
MODEL_REPO = "OpenTransformer/AGILLM-4.3"
DELTA_DIR = "checkpoints/recovery_fedC/artifacts/delta/pretrain_delta_step00374922_20260703T1248Z__sha256_f2d01389959f"
CKPT_FILE = "pretrain_delta_step00374922_20260703T1248Z.pt"
STAT_RE = re.compile(r"\[(?P<sec>[0-9.]+)s \| (?P<tok>[0-9]+) tokens \| (?P<tps>[0-9.]+) tok/s\]")
DEFAULT_PROMPT = "The quick brown fox jumps over the lazy dog and then"

print("[space] downloading checkpoint (5.4 GB, cached by HF hub)...", flush=True)
CKPT = hf_hub_download(MODEL_REPO, f"{DELTA_DIR}/{CKPT_FILE}")
TOKENIZER = hf_hub_download(MODEL_REPO, f"{DELTA_DIR}/{CKPT_FILE}.tokenizer.json")
os.environ["AGILLM43_TOKENIZER_JSON"] = TOKENIZER
print("[space] importing runtime...", flush=True)
import agillm41  # noqa: E402  (heavy single-file runtime)


class _LineWriter:
    """Captures stdout line-wise into a queue while a request runs."""

    def __init__(self, out_q):
        self.q = out_q
        self.buf = ""

    def write(self, s):
        self.buf += s
        while "\n" in self.buf:
            line, self.buf = self.buf.split("\n", 1)
            self.q.put(line)

    def flush(self):
        pass


def _build_argv(prompt, mode, max_new, nat_passes, temperature, top_k, top_p,
                greedy, ignore_eos, rep_pen, presence, frequency, last_n, streaming):
    base_mode = "sat" if str(mode).startswith("sat") else str(mode)
    argv = ["infer", "--device", "cuda", "--ckpt", CKPT, "--mode", base_mode,
            "--prompt", prompt, "--max_new", str(int(max_new)), "--min_new", "0",
            "--temperature", str(float(temperature)), "--top_k", str(int(top_k)),
            "--top_p", str(float(top_p)), "--repetition_penalty", str(float(rep_pen)),
            "--presence_penalty", str(float(presence)), "--frequency_penalty", str(float(frequency)),
            "--penalty_last_n", str(int(last_n)), "--infer_dtype", "fp16", "--plain-output"]
    if greedy:
        argv.append("--greedy")
    if ignore_eos:
        argv.append("--ignore_eos")
    if streaming:
        argv.append("--stream")
    if base_mode == "nat":
        argv.extend(["--nat_passes", str(int(nat_passes))])
    elif mode == "sat var":
        argv.append("--var")
    elif mode == "sat fixed":
        argv.append("--no-var")
    return argv


@spaces.GPU(duration=240)
def generate(prompt: str, mode: str = "nat", max_new: int = 64, nat_passes: int = 1,
             temperature: float = 0.25, top_k: int = 0, top_p: float = 1.0,
             greedy: bool = True, ignore_eos: bool = True, rep_pen: float = 2.0,
             presence: float = 0.8, frequency: float = 1.2, last_n: int = 0,
             streaming: bool = True):
    """Generate text with the AGILLM 4.3 research LLM (1.2B params, ZeroGPU).

    This function is also exposed as an MCP tool so agents can call the model.

    Args:
        prompt: Input text to continue.
        mode: Decode head — "nat" (fast parallel mask-predict), "sat var",
            "sat fixed", or "ar" (classic left-to-right). NAT is fastest.
        max_new: Number of tokens to generate.
        nat_passes: NAT refinement passes (1 is fast; more improves quality).
        temperature: Sampling temperature (ignored when greedy=True).
        top_k: Top-k filter, 0 = off.
        top_p: Nucleus sampling threshold.
        greedy: Take the argmax token each step (deterministic).
        ignore_eos: Never stop early on the end-of-sequence token.
        rep_pen: Repetition penalty (>1 discourages repeats).
        presence: Presence penalty.
        frequency: Frequency penalty.
        last_n: Penalty window in tokens, 0 = unlimited (keep 0 for NAT).
        streaming: Stream tokens as they commit.

    Returns:
        The prompt followed by the generated continuation. Note: at this
        mid-pretraining checkpoint the output is word-salad by design.
    """
    prompt = (prompt or "").strip() or DEFAULT_PROMPT
    argv = _build_argv(prompt, mode, max_new, nat_passes, temperature, top_k, top_p,
                       greedy, ignore_eos, rep_pen, presence, frequency, last_n, streaming)
    out_q = queue.Queue()
    done = threading.Event()

    def worker():
        old_argv, old_stdout = sys.argv, sys.stdout
        sys.argv = ["agillm41.py"] + argv
        sys.stdout = _LineWriter(out_q)
        try:
            agillm41.main()
        except SystemExit:
            pass
        except Exception as exc:  # surfaced to the UI below
            out_q.put(f"[SPACE_ERROR] {type(exc).__name__}: {exc}")
        finally:
            sys.stdout = old_stdout
            sys.argv = old_argv
            done.set()

    threading.Thread(target=worker, daemon=True).start()
    yield "(loading checkpoint onto the ZeroGPU slice — first request also builds the decompression cache...)"
    slots = None
    final = None
    stats = ""
    t0 = time.time()
    while not (done.is_set() and out_q.empty()):
        try:
            s = out_q.get(timeout=0.5).strip()
        except queue.Empty:
            continue
        if s.startswith("[SPACE_ERROR]"):
            final = s
            break
        if s.startswith("[STREAM_BEGIN] "):
            try:
                slots = [None] * int(json.loads(s[len("[STREAM_BEGIN] "):]).get("slots") or 0)
            except Exception:
                slots = None
            continue
        if s.startswith("[STREAM_NAT] ") or s.startswith("[STREAM_AR] "):
            if slots is None:
                continue
            try:
                d = json.loads(s.split("] ", 1)[1])
                i = int(d.get("pos", d.get("i")))
                if 0 <= i < len(slots):
                    slots[i] = str(d.get("text") or "")
            except Exception:
                pass
            yield prompt + "".join(x if x is not None else " ·" for x in slots)
            continue
        if STAT_RE.search(s):
            stats = s
            continue
        if s and not s.startswith("[") and not s.startswith("Generating"):
            final = s
    wall = time.time() - t0
    out = final or "(no output)"
    if stats:
        out += f"\n\n{stats} | wall={wall:.2f}s on ZeroGPU"
    yield out


with gr.Blocks(title="AGILLM 4.3 GUI (ZeroGPU)") as demo:
    gr.Markdown(
        "# AGILLM 4.3 — Local Inference GUI (ZeroGPU Space)\n"
        "1.2B-param research model with **AR / SAT / NAT** decode heads. NAT streaming shows the token "
        "canvas filling in confidence order — a live view of the parallel mask-predict decode. "
        "Mid-pretraining checkpoint: expect word-salad, not prose. Each request loads the model onto a "
        "ZeroGPU slice (~15-60s; the very first request also builds a disk cache and takes longer)."
    )
    with gr.Row():
        prompt = gr.Textbox(label="Prompt", value=DEFAULT_PROMPT, lines=2, scale=4)
    with gr.Row():
        mode = gr.Radio(["nat", "sat var", "sat fixed", "ar"], value="nat", label="Mode")
        streaming = gr.Checkbox(value=True, label="Streaming (live canvas)")
        greedy = gr.Checkbox(value=True, label="Greedy")
        ignore_eos = gr.Checkbox(value=True, label="Ignore EOS")
    with gr.Row():
        max_new = gr.Slider(4, 512, value=64, step=4, label="Max tokens")
        nat_passes = gr.Slider(1, 8, value=1, step=1, label="NAT passes")
        temperature = gr.Slider(0.0, 1.5, value=0.25, step=0.05, label="Temperature")
    with gr.Row():
        top_k = gr.Slider(0, 200, value=0, step=1, label="Top-k (0=off)")
        top_p = gr.Slider(0.1, 1.0, value=1.0, step=0.05, label="Top-p")
        rep_pen = gr.Slider(1.0, 4.0, value=2.0, step=0.1, label="Repetition penalty")
    with gr.Row():
        presence = gr.Slider(0.0, 2.0, value=0.8, step=0.1, label="Presence penalty")
        frequency = gr.Slider(0.0, 3.0, value=1.2, step=0.1, label="Frequency penalty")
        last_n = gr.Slider(0, 1024, value=0, step=32, label="Penalty window (0=unlimited; keep 0 for NAT)")
    out = gr.Textbox(label="Output", lines=12)
    btn = gr.Button("Run Inference", variant="primary")
    btn.click(generate,
              inputs=[prompt, mode, max_new, nat_passes, temperature, top_k, top_p,
                      greedy, ignore_eos, rep_pen, presence, frequency, last_n, streaming],
              outputs=out, api_name="generate")
    gr.Markdown(
        "**Agent / MCP access:** this Space is also an MCP server — point an MCP "
        "client at `https://opentransformer-agillm43-gui-zerogpu.hf.space/gradio_api/mcp/sse` "
        "and the `generate` tool becomes callable (each call spins up a ZeroGPU slice)."
    )

# AGILLM-MCP 20260703: mcp_server=True exposes generate() as an MCP tool.
demo.queue(default_concurrency_limit=1).launch(mcp_server=True)