"""
Hugging Face Spaces (ZeroGPU) front end for the Paper->Code grounding
pipeline -- runs Qwen3-VL-32B-Instruct-FP8 (the same model validated by
hand throughout this project) instead of Gemini, so this deploys the
project's OWN model for free, public use.

Real constraints, on purpose, not hidden:
  - ZeroGPU is Gradio-only, so this is a Gradio Blocks app, not the
    FastAPI + custom HTML used by scripts/webapp/server.py.
  - No vLLM here -- ZeroGPU allocates real GPU access only for the
    duration of an `@spaces.GPU`-decorated function call, which doesn't
    fit a long-running server process. This uses plain
    transformers.generate(), the same approach scripts/local_agent/
    run_local_agent.py already validated, just collapsed to one
    single-turn call (see run_single_turn_transformers below) instead of
    that script's multi-turn tool-calling loop.
  - Each visitor's own HF account has a small DAILY GPU-second quota
    (2 min unauthenticated, 5 min free account) -- a real single-turn
    call here is likely 70-130s, so this is genuinely "a couple of tries
    per person per day," not unlimited use.
"""
from __future__ import annotations

import base64
import io
import json
import subprocess
import sys
import time
from pathlib import Path

import gradio as gr
import spaces
import torch
from PIL import Image
from transformers import AutoModelForImageTextToText, AutoProcessor, BitsAndBytesConfig

# Locally (inside a checkout of the full PaperCodeDiagramApp repo) this
# file lives at scripts/spaces_app/app.py, three levels under the repo
# root. Uploaded to its own standalone Space repo, app.py sits at the
# Space's own root instead, alongside its own copies of the same shared
# files (RULES.md, schema/, scripts/... -- see the upload staging step).
# Detect which layout is actually on disk rather than requiring an env
# var to be configured correctly on the Space.
import os  # noqa: E402
_repo_root_override = os.environ.get("PCD_REPO_ROOT")
if _repo_root_override:
    REPO_ROOT = Path(_repo_root_override)
elif (Path(__file__).resolve().parent / "RULES.md").is_file():
    REPO_ROOT = Path(__file__).resolve().parent  # standalone Space layout
else:
    REPO_ROOT = Path(__file__).resolve().parents[2]  # local dev, full repo checkout
sys.path.insert(0, str(REPO_ROOT / "scripts" / "local_agent"))
sys.path.insert(0, str(REPO_ROOT / "scripts" / "dataset"))
sys.path.insert(0, str(REPO_ROOT / "scripts" / "webapp"))

from run_local_agent import setup_case  # noqa: E402
from run_vllm_agent import (  # noqa: E402
    RULES_PATH, SCHEMA_PATH, VERIFY_SCRIPT, _dump_narrative,
    build_system_text, collect_repo_manifest, collect_repo_source_dump,
)
from tools import CaseWorkspace, ToolError  # noqa: E402
from validation import extract_github_repo, looks_like_a_paper, normalize_paper_url, repo_exists  # noqa: E402

MODEL_ID = "Qwen/Qwen3-VL-32B-Instruct"
IMAGE_SIZE = {"shortest_edge": 16384, "longest_edge": 400000}
MAX_PAGES = 10
TOKEN_BUDGETS = [4096, 8192]  # kept modest: every extra token here is GPU time out of the visitor's own daily quota
MAX_FIX_ATTEMPTS = 1
CASES_DIR = REPO_ROOT / "cases"
UPLOADS_DIR = Path(__file__).resolve().parent / "_uploads"

print(f"loading {MODEL_ID} ...", flush=True)
processor = AutoProcessor.from_pretrained(MODEL_ID, size=IMAGE_SIZE)
# 8-bit bitsandbytes instead of the FP8 checkpoint: a bf16 32B model alone
# needs ~64GB (too much for ZeroGPU's default 48GB budget), and the FP8
# checkpoint's quantized kernels are downloaded from the Hub at runtime by
# the `kernels` package, which repeatedly failed to find a build matching
# this container's torch/CUDA. bitsandbytes' kernels ship inside the pip
# package itself -- no runtime download -- and int8 needs ~32GB, same
# budget the FP8 checkpoint was chosen for.
model = AutoModelForImageTextToText.from_pretrained(
    MODEL_ID,
    quantization_config=BitsAndBytesConfig(load_in_8bit=True),
    device_map="cuda",
)
model.eval()
print("model loaded", flush=True)


def _b64_to_pil(data: str) -> Image.Image:
    return Image.open(io.BytesIO(base64.b64decode(data))).convert("RGB")


def _build_messages(ws: CaseWorkspace, system_text: str, schema: dict) -> list[dict]:
    content: list[dict] = [{"type": "text", "text": (
        "SINGLE-TURN TASK: you get exactly this one message and must produce "
        "exactly one response -- the complete, final narrative.json object as "
        "raw JSON, nothing else (no prose, no markdown fences, no tool calls -- "
        "there are no tools in this mode). Every PDF page you need and the "
        "repo's full source are already attached below; you will not get a "
        "chance to ask for anything else or to revise after this. Skip figures "
        "entirely: omit `case.overview.figure` and every node's `figure_refs`. "
        "Ground every code_ref in the attached source using its real file path "
        "and line numbers (the numbers prefixing each line below). A full "
        "manifest of every file's REAL path and REAL line count is included "
        "below the page images -- check every citation's file path and line "
        "range against that manifest before writing it; never cite a path "
        "that isn't in the manifest, and never cite a line past a file's "
        "listed length. Make sure start_line <= end_line for every citation. "
        "Respond with ONLY the raw JSON object, starting with `{` and ending "
        f"with `}}`. The schema your JSON must match:\n{json.dumps(schema)}"
    )}]
    for page in range(1, MAX_PAGES + 1):
        try:
            blocks = ws.read_pdf_page(page)
        except ToolError:
            break
        content.append({"type": "text", "text": f"--- paper.pdf page {page} ---"})
        for b in blocks:
            if b["type"] == "image":
                content.append({"type": "image", "image": _b64_to_pil(b["source"]["data"])})

    manifest = collect_repo_manifest(ws)
    content.append({"type": "text", "text": f"--- repo/ file manifest (real path : real line count) ---\n{manifest}"})
    repo_dump = collect_repo_source_dump(ws)
    content.append({"type": "text", "text": f"--- full repo source (repo/) ---\n{repo_dump}"})

    return [
        {"role": "system", "content": system_text},
        {"role": "user", "content": content},
    ]


def _collect_images(messages: list[dict]) -> list[Image.Image]:
    imgs = []
    for m in messages:
        content = m.get("content")
        if isinstance(content, list):
            for item in content:
                if isinstance(item, dict) and item.get("type") == "image":
                    imgs.append(item["image"])
    return imgs


def _extract_json(raw: str) -> dict | None:
    """The model was asked for raw JSON only, but strip a markdown fence
    if it added one anyway -- cheap and doesn't mask a real parse failure
    for anything else."""
    text = raw.strip()
    if text.startswith("```"):
        text = text.split("\n", 1)[1] if "\n" in text else text
        if text.endswith("```"):
            text = text.rsplit("```", 1)[0]
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        return None


@spaces.GPU(duration=90)
def _generate(messages: list[dict], max_new_tokens: int) -> tuple[str, str]:
    """The only part that actually needs the GPU. Returns (generated_text,
    finish_reason) -- finish_reason is "length" if max_new_tokens was hit
    (transformers doesn't expose a stop-reason the way an API does, so
    this is inferred: did decoding stop because of an EOS token, or
    because we ran out of budget)."""
    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    images = _collect_images(messages)
    inputs = processor(text=[text], images=[images] if images else None, return_tensors="pt").to(model.device)
    with torch.no_grad():
        out_ids = model.generate(
            **inputs, max_new_tokens=max_new_tokens, do_sample=True,
            temperature=0.4, top_p=0.9, repetition_penalty=1.1,
        )
    new_ids = out_ids[0][inputs["input_ids"].shape[1]:]
    finish_reason = "length" if len(new_ids) >= max_new_tokens else "stop"
    generated = processor.tokenizer.decode(new_ids, skip_special_tokens=True)
    return generated, finish_reason


def _generate_narrative_once(messages: list[dict], log) -> tuple[dict | None, float]:
    total_latency = 0.0
    for attempt, budget in enumerate(TOKEN_BUDGETS, start=1):
        t0 = time.perf_counter()
        raw, finish_reason = _generate(messages, budget)
        latency = time.perf_counter() - t0
        total_latency += latency
        log(f"--- generation attempt {attempt}/{len(TOKEN_BUDGETS)} "
            f"(max_new_tokens={budget}, {latency:.1f}s, finish_reason={finish_reason}) ---")
        if finish_reason == "length":
            log(f"truncated at {budget} tokens" + (" -- retrying with more room" if attempt < len(TOKEN_BUDGETS) else " -- out of retries"))
            continue
        narrative = _extract_json(raw)
        if narrative is not None:
            return narrative, total_latency
        log("FAILED: response was not valid JSON" + (" -- retrying" if attempt < len(TOKEN_BUDGETS) else " -- out of retries"))
    return None, total_latency


def ground_paper(paper_url: str, progress=gr.Progress()):
    log_lines: list[str] = []

    def log(msg: str) -> None:
        log_lines.append(msg)

    if not paper_url or not paper_url.strip():
        return "Enter a paper URL first.", None, gr.update()

    progress(0.05, desc="downloading + validating the paper")
    url = normalize_paper_url(paper_url.strip())
    UPLOADS_DIR.mkdir(parents=True, exist_ok=True)
    staged_pdf = UPLOADS_DIR / f"{abs(hash(url))}.pdf"
    try:
        import urllib.request
        urllib.request.urlretrieve(url, staged_pdf)
    except Exception as e:  # noqa: BLE001
        return f"Couldn't download {url}: {e}", None, gr.update()

    ok, reason = looks_like_a_paper(staged_pdf)
    if not ok:
        staged_pdf.unlink(missing_ok=True)
        return f"(1) doesn't look like a paper: {reason}", None, gr.update()

    repo_url = extract_github_repo(staged_pdf)
    if not repo_url:
        staged_pdf.unlink(missing_ok=True)
        return "(2) no GitHub repository link found anywhere in this paper's text", None, gr.update()

    progress(0.15, desc=f"found repo: {repo_url} -- cloning")
    import re
    slug = re.sub(r"[^a-z0-9]+", "-", staged_pdf.stem.lower()).strip("-") or "case"
    case_dir, slug = setup_case(str(staged_pdf), repo_url, CASES_DIR, slug)

    if not repo_exists(repo_url):
        return f"Repo not reachable: {repo_url}", None, gr.update()

    progress(0.3, desc="reading paper + repo, building the prompt")
    ws = CaseWorkspace(case_dir)
    system_text = build_system_text(None, None)
    schema = json.loads(SCHEMA_PATH.read_text())
    messages = _build_messages(ws, system_text, schema)

    progress(0.4, desc="calling the model (this is the slow part)")
    narrative, latency = _generate_narrative_once(messages, log)
    if narrative is None:
        return "The model didn't produce a valid response:\n\n" + "\n".join(log_lines), None, gr.update()

    fix_attempts = 0
    while True:
        progress(0.8, desc="verifying citations against the real code")
        resolved = ws.resolve_code_snippets(narrative)
        narrative_path = case_dir / "narrative.json"
        narrative_path.write_text(json.dumps(resolved, indent=2, ensure_ascii=False) + "\n")
        verify = subprocess.run(
            [sys.executable, str(VERIFY_SCRIPT), str(narrative_path)],
            capture_output=True, text=True,
        )
        log(f"=== verify_case.py (returncode={verify.returncode}) ===\n{verify.stdout}\n{verify.stderr}")

        if verify.returncode == 0:
            _dump_narrative(resolved)
            progress(1.0, desc="done")
            return _summary_markdown(resolved, True, latency, log_lines), resolved, gr.update(choices=_list_case_choices(), value=slug)

        if fix_attempts >= MAX_FIX_ATTEMPTS:
            _dump_narrative(resolved)
            return _summary_markdown(resolved, False, latency, log_lines), resolved, gr.update(choices=_list_case_choices(), value=slug)

        fix_attempts += 1
        progress(0.6, desc=f"fixing a citation issue (attempt {fix_attempts}/{MAX_FIX_ATTEMPTS})")
        messages.append({"role": "assistant", "content": json.dumps(narrative)})
        messages.append({"role": "user", "content": (
            f"scripts/verify_case.py rejected that submission:\n\n{verify.stdout}\n{verify.stderr}\n\n"
            "This is your only chance to fix it. Re-check the file manifest above "
            "for the real path/line count of anything flagged, then resubmit the "
            "COMPLETE corrected narrative.json as raw JSON (same format as "
            "before, not a diff) -- fix exactly the issues named above and "
            "nothing else."
        )})
        narrative, gen_latency = _generate_narrative_once(messages, log)
        latency += gen_latency
        if narrative is None:
            return "Fix attempt didn't produce a valid response:\n\n" + "\n".join(log_lines), None, gr.update()


def _summary_markdown(narrative: dict, ok: bool, latency: float, log_lines: list[str]) -> str:
    case = narrative.get("case", {})
    nodes = narrative.get("nodes", [])
    lines = [
        f"## {case.get('paper_title', '(untitled)')}",
        f"*{case.get('venue', '')}* &middot; `{case.get('repo', '')}`",
        "",
        "**Verified: yes** — every citation checked against the real code." if ok else
        "**Verifier found an issue** — see the log below; treat citations with care.",
        f"Model call time: {latency:.1f}s",
        "",
        "### Overview",
    ]
    task = (case.get("overview") or {}).get("task", {})
    lines.append(f"- **Input:** {task.get('input', '')}")
    lines.append(f"- **Task:** {task.get('task', '')}")
    lines.append(f"- **Output:** {task.get('output', '')}")
    lines.append("")
    lines.append("### Contributions")
    for n in nodes:
        lines.append(f"**{n.get('title')}** ({n.get('status')}) &mdash; {n.get('why', '')}")
        for cr in n.get("code_refs", []):
            lines.append(f"  - `{cr.get('file')}:{cr.get('start_line')}-{cr.get('end_line')}`")
    results = case.get("results", [])
    if results:
        lines.append("")
        lines.append("### Results")
        for t in results:
            lines.append(f"**{t.get('table_label', '')}** &mdash; {t.get('caption', '')}")
            cols = t.get("columns", [])
            if cols:
                lines.append("| " + " | ".join(cols) + " |")
                lines.append("|" + "---|" * len(cols))
                for row in t.get("rows", []):
                    lines.append("| " + " | ".join(row) + " |")
    lines.append("")
    lines.append("<details><summary>Full run log</summary>\n\n```\n" + "\n".join(log_lines)[-4000:] + "\n```\n</details>")
    return "\n".join(lines)


def _list_case_choices() -> list[str]:
    if not CASES_DIR.is_dir():
        return []
    return sorted(p.name for p in CASES_DIR.iterdir() if (p / "narrative.json").is_file())


def load_past_case(slug: str):
    if not slug:
        return "", None
    narrative_path = CASES_DIR / slug / "narrative.json"
    if not narrative_path.is_file():
        return "No such case.", None
    narrative = json.loads(narrative_path.read_text())
    return _summary_markdown(narrative, True, 0.0, ["(loaded from disk, not re-run)"]), narrative


with gr.Blocks(title="Paper -> Code Grounding (Qwen3-VL, free ZeroGPU)") as demo:
    gr.Markdown(
        "# Paper -> Code Grounding\n"
        "Runs **Qwen3-VL-32B-Instruct (FP8)** on Hugging Face's free ZeroGPU. "
        "Only proceeds if the paper is real and links a GitHub repo.\n\n"
        "**Heads up on quota**: ZeroGPU gives every visitor's own HF account a small daily GPU budget "
        "(2 min if you're not logged in, 5 min on a free account) -- a real run here takes roughly a "
        "minute or two, so expect a handful of tries per day, not unlimited use."
    )
    with gr.Row():
        url_in = gr.Textbox(label="Paper URL (arXiv link)", placeholder="https://arxiv.org/abs/XXXX.XXXXX", scale=4)
        submit_btn = gr.Button("Ground it", variant="primary", scale=1)
    with gr.Row():
        past_cases = gr.Dropdown(label="Past cases", choices=_list_case_choices(), interactive=True)
    output_md = gr.Markdown()
    narrative_state = gr.JSON(visible=False)

    submit_btn.click(ground_paper, inputs=[url_in], outputs=[output_md, narrative_state, past_cases])
    past_cases.change(load_past_case, inputs=[past_cases], outputs=[output_md, narrative_state])

if __name__ == "__main__":
    demo.queue().launch()