"""Pulpie vs Dripper — encoder vs decoder content extraction, live.

gradio.Server backend: a FastAPI app with Gradio's queue engine layered on
(via @app.api()). The UI is a custom vanilla HTML/CSS/JS frontend served at
`/`; the Gradio JS client calls the queued `/race` endpoint, so requests go
through concurrency control and (on ZeroGPU Spaces) GPU allocation — not raw
fetches that would collide on the GPU.

Both models receive the SAME simplified HTML (pulpie's MinerU-HTML-parity
simplify). Pulpie classifies every block in one encoder forward pass; Dripper
(MinerU-HTML 0.6B) emits its answer autoregressively, one token at a time.
"""

from __future__ import annotations

import os
import re
import threading
import time
from pathlib import Path

import spaces
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
from fastapi.responses import HTMLResponse
from gradio import Server
from gradio.data_classes import FileData

from pulpie.chunker import extract_blocks, pack_chunks, tokenize_blocks
from pulpie.markdown import to_markdown
from pulpie.model_utils import (
    extract_item_ids,
    load_model_and_tokenizer,
    predictions_to_labels,
    resolve_model_id,
)
from pulpie.reconstruct import extract_main_html
from pulpie.simplify import simplify

PULPIE_MODEL = "feyninc/pulpie-orange-small"
DRIPPER_MODEL = "opendatalab/MinerU-HTML-v1.1-hunyuan0.5B-compact"
# Dripper's real context is 32768 (matches the benchmark harness setting);
# leave headroom for the answer tokens.
DRIPPER_MAX_INPUT_TOKENS = 32768 - 8192
DRIPPER_MAX_NEW_TOKENS = 8192
MAX_HTML_BYTES = 2_000_000

EXAMPLES_DIR = Path(__file__).parent / "examples"
EXAMPLES = {
    "paulgraham.com": "01_article_pg.html",
    "nodejs.dev": "02_nodejs_dev.html",
    "protobuf.dev": "05_protobuf_dev.html",
    "istio.io": "09_istio_io.html",
    "tutorialspoint.com": "11_tutorialspoint_com.html",
    "builtin.com": "12_builtin_com.html",
    "bbc.com": "15_bbc_com.html",
    "theguardian.com": "17_theguardian_com.html",
    "wikipedia.org": "14_wikipedia_org.html",
    "stackoverflow.com": "16_stackoverflow_com.html",
}

# MinerU-HTML's own "short_compact" prompt (build_prompt.py, Apache-2.0) —
# the format Dripper was trained on.
DRIPPER_PROMPT = """As an HTML expert, classify elements with "_item_id" as "main" or "other", keeping only the main content and removing nav, metadata, etc.
Guidelines:
"Main": Includes primary content like article text, images in the article, original posts in forums, Q&A questions, and answers.
"Other": Includes navigation, metadata, ads, user info, and related content (e.g., sidebars, timestamps, suggested articles).
Output Format:
Return each _item_id and its corresponding category in the following format:
{{_item_id}}{{category}}{{_item_id}}{{category}}...
Here, {{_item_id}} may be 1, 2, 3..., and {{category}} is either "main" or "other", as shown in the example below:
"1main2other3other4main"
Input HTML:
{html}"""

# ── Model loading (CPU at startup; moved to CUDA inside @spaces.GPU) ──

# ZeroGPU pattern: .to("cuda") at startup — spaces defers the actual move to
# the GPU slice, so per-call transfer cost disappears.
_STARTUP_DEVICE = "cuda" if (spaces.config.Config.zero_gpu or torch.cuda.is_available()) else "cpu"

print("Loading pulpie orange-small...")
_pulpie_model, _pulpie_tokenizer, _sep_token_id = load_model_and_tokenizer(
    resolve_model_id(PULPIE_MODEL), torch.device("cpu")
)
_pulpie_model = _pulpie_model.to(_STARTUP_DEVICE)
_pulpie_pad_id = _pulpie_model.config.pad_token_id or 0

print("Loading Dripper...")
_dripper_tokenizer = AutoTokenizer.from_pretrained(DRIPPER_MODEL, trust_remote_code=True)
_dripper_model = AutoModelForCausalLM.from_pretrained(
    DRIPPER_MODEL, dtype=torch.bfloat16, trust_remote_code=True
).to(_STARTUP_DEVICE)


@spaces.GPU(duration=30)
def pulpie_classify(simplified: str) -> tuple[dict[str, str], float]:
    """One (or few) encoder forward passes → (block labels, inference seconds)."""
    device = torch.device(_STARTUP_DEVICE)
    model = _pulpie_model

    blocks = extract_blocks(simplified)
    if not blocks:
        return {}, 0.0
    item_ids = extract_item_ids(blocks)
    block_token_ids = tokenize_blocks(blocks, _pulpie_tokenizer)
    chunks = pack_chunks(
        block_token_ids,
        max_tokens=8192,
        sep_token_id=_sep_token_id,
        bos_token_id=_pulpie_tokenizer.bos_token_id,
        eos_token_id=_pulpie_tokenizer.eos_token_id,
    )

    t0 = time.perf_counter()
    predictions = [0] * len(blocks)
    with torch.no_grad():
        for chunk_ids, block_indices in chunks:
            input_ids = torch.tensor([chunk_ids], dtype=torch.long, device=device)
            attention_mask = torch.ones_like(input_ids)
            logits = model(input_ids=input_ids, attention_mask=attention_mask).logits[0]
            sep_positions = (input_ids[0] == _sep_token_id).nonzero(as_tuple=True)[0]
            preds = logits[sep_positions].argmax(dim=-1).cpu().tolist()
            for idx, block_idx in enumerate(block_indices):
                if idx < len(preds):
                    predictions[block_idx] = preds[idx]
    if device.type == "cuda":
        torch.cuda.synchronize()

    return predictions_to_labels(item_ids, predictions), time.perf_counter() - t0


@spaces.GPU(duration=120)
def dripper_generate_stream(simplified: str):
    """Decoder generation, yielding (accumulated_text, inference_seconds, error)."""
    device = torch.device(_STARTUP_DEVICE)
    model = _dripper_model

    prompt = DRIPPER_PROMPT.format(html=simplified)
    messages = [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": prompt},
    ]
    chat_prompt = _dripper_tokenizer.apply_chat_template(
        messages, tokenize=False, enable_thinking=False, add_generation_prompt=True
    )
    inputs = _dripper_tokenizer(chat_prompt, return_tensors="pt")
    inputs.pop("token_type_ids", None)  # tokenizer emits it; the causal LM rejects it
    if inputs.input_ids.shape[1] > DRIPPER_MAX_INPUT_TOKENS:
        yield None, 0.0, "PAGE TOO LONG FOR DRIPPER'S CONTEXT — a failure mode pulpie avoids."
        return
    inputs = {k: v.to(device) for k, v in inputs.items()}

    streamer = TextIteratorStreamer(
        _dripper_tokenizer, skip_prompt=True, skip_special_tokens=True
    )
    thread = threading.Thread(
        target=model.generate,
        kwargs=dict(
            **inputs,
            max_new_tokens=DRIPPER_MAX_NEW_TOKENS,
            do_sample=False,
            streamer=streamer,
        ),
    )
    t0 = time.perf_counter()
    thread.start()
    acc = ""
    for piece in streamer:
        acc += piece
        yield acc, time.perf_counter() - t0, None
    thread.join()
    yield acc, time.perf_counter() - t0, None


def parse_dripper_labels(text: str) -> dict[str, str]:
    return {m.group(1): m.group(2) for m in re.finditer(r"(\d+)(main|other)", text)}


def labels_to_markdown(map_html: str, labels: dict[str, str]) -> str:
    if not labels or not any(v == "main" for v in labels.values()):
        return ""
    return to_markdown(extract_main_html(map_html, labels))


# ── The race (queued endpoint, streamed as SSE via the Gradio client) ──

MONO = "font-family: 'IBM Plex Mono', monospace;"


def fmt_timer(seconds: float, done: bool, color: str = "#1B1714") -> str:
    state = "" if done else '<span style="color:#2F5A3C;">● RUNNING</span> '
    return (
        f'<div style="{MONO} font-size:13px; letter-spacing:0.08em; color:{color};">'
        f"{state}{seconds:.1f}s</div>"
    )


def race(example_name: str, custom_html: str):
    """Generator driving both panes. Yields dict records consumed by the SSE
    frontend: {pulpie_md, pulpie_timer, dripper_md, dripper_timer, verdict}."""
    if custom_html and custom_html.strip():
        html = custom_html
    elif example_name in EXAMPLES:
        html = (EXAMPLES_DIR / EXAMPLES[example_name]).read_text(errors="replace")
    else:
        yield {"pulpie_md": "", "pulpie_timer": "", "dripper_md": "",
               "dripper_timer": "", "verdict": "Pick an example page or paste HTML."}
        return
    if len(html.encode()) > MAX_HTML_BYTES:
        yield {"pulpie_md": "", "pulpie_timer": "", "dripper_md": "",
               "dripper_timer": "", "verdict": "HTML too large (2 MB cap)."}
        return

    # Shared preprocessing — identical input to both models.
    try:
        simplified, map_html = simplify(html)
    except Exception as e:
        yield {"pulpie_md": "", "pulpie_timer": "", "dripper_md": "",
               "dripper_timer": "", "verdict": f"simplify failed: {e}"}
        return

    # Pulpie. Timed on-GPU inside the @spaces.GPU call so ZeroGPU
    # slice-attach overhead (identical for both models) isn't counted.
    pulpie_labels, pulpie_s = pulpie_classify(simplified)
    pulpie_md = labels_to_markdown(map_html, pulpie_labels) or "*No main content detected.*"
    pulpie_timer = fmt_timer(pulpie_s, done=True, color="#C6531D")
    yield {"pulpie_md": pulpie_md, "pulpie_timer": pulpie_timer,
           "dripper_md": "", "dripper_timer": fmt_timer(0.0, done=False), "verdict": ""}

    # Dripper, streaming.
    raw, dripper_s = "", 0.0
    for acc, gen_s, err in dripper_generate_stream(simplified):
        if err:
            yield {"pulpie_md": pulpie_md, "pulpie_timer": pulpie_timer,
                   "dripper_md": f"*{err}*", "dripper_timer": fmt_timer(gen_s, done=True),
                   "verdict": ""}
            return
        raw, dripper_s = acc, gen_s
        preview = f"```\n…{raw[-1500:]}\n```" if raw else "*waiting for first token…*"
        yield {"pulpie_md": pulpie_md, "pulpie_timer": pulpie_timer,
               "dripper_md": preview, "dripper_timer": fmt_timer(dripper_s, done=False),
               "verdict": ""}

    dripper_md = labels_to_markdown(map_html, parse_dripper_labels(raw)) or "*No main content detected.*"
    speedup = dripper_s / max(pulpie_s, 1e-6)
    speedup_str = f"{speedup:.1f}x" if speedup < 10 else f"{speedup:.0f}x"
    verdict = (
        f'<div style="{MONO} font-size:14px; letter-spacing:0.08em; padding:14px 0; '
        f'border-top:1px solid #D7CBB6; color:#1B1714;">'
        f"PULPIE {pulpie_s:.2f}s → DRIPPER {dripper_s:.1f}s · "
        f'<span style="color:#C6531D; font-weight:600;">{speedup_str} ON THIS PAGE</span></div>'
    )
    yield {"pulpie_md": pulpie_md, "pulpie_timer": pulpie_timer,
           "dripper_md": dripper_md, "dripper_timer": fmt_timer(dripper_s, done=True),
           "verdict": verdict}


# ── gradio.Server app ──

app = Server()


@app.api(name="race")
def race_endpoint(example_name: str, custom_html: str) -> dict:
    """Queued, GPU-managed streaming endpoint. Each yielded record is delivered
    to the Gradio client (JS or Python) as a separate SSE event, so the frontend
    sees pulpie finish, then dripper's tokens arrive live, then the verdict.
    The final record carries every field, so a non-streaming predict() call
    still converges to the complete result.

    The `-> dict` return annotation is load-bearing: gradio.Server derives the
    endpoint's output components from the return type, and a generator without
    one registers zero outputs — the queue runs (the logs show each yield) but
    the client receives nothing to render. Annotating with the *yielded* type
    (dict, not Generator[dict, ...]) declares one dict output while the queue
    still iterates the generator and streams each dict as a data event."""
    yield from race(example_name or "", custom_html or "")


@app.get("/")
async def homepage():
    html_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "index.html")
    with open(html_path, "r", encoding="utf-8") as f:
        return HTMLResponse(content=f.read())


if __name__ == "__main__":
    app.launch(show_error=True)