"""Dialectical abstracts demo, full-bleed UI.

Left pane: the paper (abstract withheld). Right rail: intuition cards with
agree / pass verdicts; agreed cards rise to the top, passed ones shelve at the
bottom; a generate button streams the model's thinking, then the abstract,
then per-intuition self-score bars.

Architecture (validated in the dialectical-intuitions Space): custom HTML UI
drives hidden gradio components via the native-setter bridge, so a browser
visitor's own ZeroGPU quota is used; streaming reaches the page through a
CPU-side peek endpoint keyed by a client job id.
"""
import json
import os
import re
import threading
from collections import OrderedDict

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

VER = "c1"
ADAPTER_REPO = os.environ.get("ADAPTER_REPO", "andreiski/dialectical-abstracts-sft-v1")
BASE = "Qwen/Qwen3-8B"
HF_TOKEN = os.environ.get("HF_TOKEN")

PAPERS = json.load(open("demo_papers.json"))
DECKS = {}
if os.path.exists("demo_decks.jsonl"):
    with open("demo_decks.jsonl") as f:
        for line in f:
            d = json.loads(line)
            DECKS[d["paper_id"]] = d
TMPL = open("writer_prompt.txt").read()

STATE = OrderedDict()          # jid -> latest chunk json (peek side channel)
MAX_STATE = 40

_model = None
_tok = None


def _load():
    global _model, _tok
    if _model is None:
        _tok = AutoTokenizer.from_pretrained(BASE)
        base = AutoModelForCausalLM.from_pretrained(
            BASE, torch_dtype=torch.bfloat16, device_map="cuda")
        from peft import PeftModel
        m = PeftModel.from_pretrained(base, ADAPTER_REPO, token=HF_TOKEN)
        _model = m.merge_and_unload()
        _model.eval()
    return _model, _tok


def clean_title(t):
    return re.sub(r"^\[[^\]]+\]\s*", "", t or "").strip()


def body_of(rec):
    parts = []
    for s in rec["sections"]:
        parts.append(s["title"])
        parts.extend(s["paras"])
    return "\n".join(parts)


def split_out(acc):
    if "</think>" in acc:
        thinking, _, rest = acc.partition("</think>")
        thinking = thinking.replace("<think>", "").strip()
    else:
        thinking, rest = acc.replace("<think>", "").strip(), ""
    m = re.search(r"(?mi)^\s*\**\s*SELF-ASSESSMENT\**:?", rest)
    abstract = rest[:m.start()].strip() if m else rest.strip()
    abstract = re.sub(r"^\s*\*{0,2}\s*Abstract:?\s*\*{0,2}\s*\n+", "", abstract, flags=re.I).strip()
    sa = rest[m.end():] if m else ""
    scores = {}
    for line in sa.splitlines():
        mm = re.match(r"\s*(\d+)\s*[:.]\s*([1-5])", line)
        if mm:
            scores[int(mm.group(1))] = int(mm.group(2))
    return thinking, abstract, scores


@spaces.GPU(duration=120)
def _gpu_generate(prompt):
    model, tok = _load()
    text = tok.apply_chat_template(
        [{"role": "user", "content": prompt}], tokenize=False,
        add_generation_prompt=True, enable_thinking=True)
    ids = tok(text, return_tensors="pt", add_special_tokens=False).input_ids.cuda()
    streamer = TextIteratorStreamer(tok, skip_prompt=True, skip_special_tokens=True)
    kw = dict(input_ids=ids, max_new_tokens=2600, do_sample=True,
              temperature=0.8, top_p=0.95, streamer=streamer)
    t = threading.Thread(target=model.generate, kwargs=kw)
    t.start()
    acc = ""
    for chunk in streamer:
        acc += chunk
        yield acc
    t.join()
    yield acc


def generate(payload):
    """Plain generator wrapper; each chunk stamped with the client jid and
    mirrored to STATE so the page can poll peek() independently."""
    try:
        req = json.loads(payload)
        jid = req["jid"]
        pid = req["paper_id"]
        sel = req["intuitions"]
    except Exception:
        yield json.dumps({"phase": "error", "msg": "bad request"})
        return
    rec = next((p for p in PAPERS if p["id"] == pid), None)
    if rec is None or not sel:
        out = json.dumps({"jid": jid, "phase": "error", "msg": "pick a paper and hold at least one intuition"})
        STATE[jid] = out
        yield out
        return
    prompt = TMPL.format(
        intuitions="\n".join(f"{i+1}. {t}" for i, t in enumerate(sel)),
        body=body_of(rec))
    final = ""
    for acc in _gpu_generate(prompt):
        final = acc
        thinking, abstract, scores = split_out(acc)
        out = json.dumps({"jid": jid, "phase": "run", "thinking": thinking,
                          "abstract": abstract})
        STATE[jid] = out
        while len(STATE) > MAX_STATE:
            STATE.popitem(last=False)
        yield out
    thinking, abstract, scores = split_out(final)
    out = json.dumps({"jid": jid, "phase": "done", "thinking": thinking,
                      "abstract": abstract, "scores": scores})
    STATE[jid] = out
    yield out


def peek(jid):
    return STATE.get(jid, "")


def ping():
    return "pong"


def papers_payload():
    return PAPERS_JS


PAPERS_JS = json.dumps([
    {"id": p["id"], "title": clean_title(p["title"]),
     "body": [{"t": s["title"], "p": s["paras"]} for s in p["sections"]],
     "rhetorical": DECKS.get(p["id"], {}).get("rhetorical", []),
     "content": DECKS.get(p["id"], {}).get("content", [])}
    for p in PAPERS])

UI_HTML = """
<div id="ivx-app">
  <section id="ivx-left">
    <div id="ivx-lhead">
      <span id="ivx-ptitle"></span>
      <span id="ivx-status" hidden></span>
      <span id="ivx-ver">VERSION</span>
    </div>
    <div id="ivx-empty">Hold a few intuitions on the right, then press write.</div>
    <div id="ivx-out" hidden>
      <details id="ivx-thinkbox" open>
        <summary>model thinking</summary>
        <div id="ivx-think"></div>
      </details>
      <div id="ivx-abs"></div>
      <div id="ivx-bars"></div>
    </div>
  </section>
  <aside id="ivx-rail">
    <div id="ivx-tabs"></div>
    <div id="ivx-pilewrap">
      <div id="ivx-counter"></div>
      <div id="ivx-pile">
        <div id="ivx-card"><span class="kind"></span><span class="txt"></span></div>
      </div>
      <div id="ivx-verdicts">
        <button id="ivx-pass" title="pass (left arrow)">&#10005;</button>
        <button id="ivx-hold" title="hold (right arrow)">&#10003;</button>
      </div>
    </div>
    <button id="ivx-gen" disabled>write the abstract</button>
    <div id="ivx-heldwrap">
      <div id="ivx-heldhead" hidden>held, newest first</div>
      <div id="ivx-held"></div>
    </div>
  </aside>
</div>
"""

UI_CSS = """
#ivx-app, #ivx-app * { box-sizing: border-box; }
#ivx-app {
  --bg:#eceef1; --surface:#ffffff; --ink:#23262d; --muted:#6a7078; --line:#d8dbe0;
  --accent:#2358b8; --hold:#1f7a4d; --hold-bg:#e9f4ee; --pass-ink:#8b939c;
  --think-bg:#f4f5f7;
  display:grid; grid-template-columns:minmax(0,1.1fr) minmax(400px,.9fr);
  height:96vh; margin:0; background:var(--bg); color:var(--ink);
  font-family:system-ui,-apple-system,"Segoe UI",Roboto,sans-serif; line-height:1.55;
  border:1px solid var(--line); border-radius:10px; overflow:hidden;
}
@media (prefers-color-scheme: dark) {
  #ivx-app { --bg:#141619; --surface:#1d2026; --ink:#e7e5e0; --muted:#9aa0a8;
    --line:#31353c; --accent:#8fb0f2; --hold:#6fce9e; --hold-bg:#1d2b24;
    --pass-ink:#77808a; --think-bg:#181b20; }
}
#ivx-left { overflow-y:auto; background:var(--surface); border-right:1px solid var(--line);
  padding:22px 28px 50px; }
#ivx-lhead { display:flex; align-items:baseline; gap:12px; margin-bottom:14px; }
#ivx-ptitle { font-weight:700; font-size:16px; flex:1; }
#ivx-status { font-size:12.5px; color:var(--muted); }
#ivx-status[hidden] { display:none; }
#ivx-ver { font-size:11px; color:var(--muted); }
#ivx-empty { color:var(--muted); font-size:14px; margin-top:30vh; text-align:center; }
#ivx-out[hidden] { display:none; }
#ivx-thinkbox { background:var(--think-bg); border:1px solid var(--line); border-radius:8px;
  padding:8px 12px; margin-bottom:14px; }
#ivx-thinkbox summary { font-size:12px; color:var(--muted); cursor:pointer; }
#ivx-think { font-size:12.5px; color:var(--muted); white-space:pre-wrap; max-height:38vh;
  overflow-y:auto; margin-top:8px; }
#ivx-abs { font-family:"Iowan Old Style",Palatino,Georgia,serif; font-size:16px;
  line-height:1.65; margin-bottom:16px; }
#ivx-abs:empty { display:none; }
.barrow { display:flex; align-items:center; gap:8px; margin:4px 0; }
.bartext { flex:1; font-size:12px; color:var(--muted); overflow:hidden; text-overflow:ellipsis;
  white-space:nowrap; }
.bartrack { width:130px; height:8px; background:var(--line); border-radius:4px; overflow:hidden; }
.barfill { height:100%; }
.barnum { width:30px; font-size:11px; color:var(--muted); text-align:right; }
#ivx-rail { overflow-y:auto; padding:16px; display:flex; flex-direction:column; gap:12px; }
#ivx-tabs { display:flex; gap:6px; flex-wrap:wrap; }
#ivx-tabs button { border:1px solid var(--line); background:var(--surface); color:var(--muted);
  font:inherit; font-size:12px; padding:5px 11px; border-radius:999px; cursor:pointer; }
#ivx-tabs button.on { background:var(--ink); color:var(--bg); border-color:var(--ink); }
#ivx-pilewrap { display:flex; flex-direction:column; align-items:center; gap:10px; }
#ivx-counter { font-size:11.5px; color:var(--muted); letter-spacing:.04em; }
#ivx-pile { position:relative; width:100%; min-height:150px; }
#ivx-pile::before, #ivx-pile::after { content:""; position:absolute; inset:0;
  background:var(--surface); border:1px solid var(--line); border-radius:12px; z-index:0; }
#ivx-pile::before { transform:translateY(8px) scale(.96); opacity:.6; }
#ivx-pile::after { transform:translateY(4px) scale(.98); opacity:.8; }
#ivx-card { position:relative; z-index:2; background:var(--surface); border:1px solid var(--line);
  border-radius:12px; padding:18px 18px 20px; min-height:150px; display:flex;
  flex-direction:column; gap:8px; box-shadow:0 6px 18px rgba(15,20,30,.08);
  transition:transform .28s ease, opacity .28s ease; }
#ivx-card .kind { font-size:10px; letter-spacing:.08em; text-transform:uppercase;
  color:var(--muted); }
#ivx-card .txt { font-size:14.5px; }
#ivx-card.fly-left { transform:translateX(-120%) rotate(-6deg); opacity:0; }
#ivx-card.fly-right { transform:translateX(120%) rotate(6deg); opacity:0; }
#ivx-verdicts { display:flex; gap:26px; }
#ivx-verdicts button { width:52px; height:52px; border-radius:50%; border:1.5px solid var(--line);
  background:var(--surface); font-size:20px; cursor:pointer; line-height:1;
  transition:transform .12s; }
#ivx-verdicts button:hover { transform:scale(1.08); }
#ivx-pass { color:var(--pass-ink); }
#ivx-hold { color:var(--hold); }
#ivx-gen { border:0; background:var(--accent); color:#fff; font:inherit; font-size:14.5px;
  font-weight:600; padding:10px 16px; border-radius:8px; cursor:pointer; }
#ivx-gen:disabled { opacity:.45; cursor:default; }
#ivx-heldhead { font-size:11px; letter-spacing:.1em; text-transform:uppercase;
  color:var(--muted); margin-bottom:6px; }
#ivx-heldhead[hidden] { display:none; }
.hcard { background:var(--hold-bg); border:1px solid var(--line); border-left:4px solid var(--hold);
  border-radius:8px; padding:8px 10px; margin-bottom:7px; display:flex; gap:8px;
  align-items:flex-start; font-size:12.5px; animation:drop .25s ease; }
@keyframes drop { from { transform:translateY(-8px); opacity:0; } to { transform:none; opacity:1; } }
.hcard .txt { flex:1; }
.hcard .rm { border:0; background:none; color:var(--muted); cursor:pointer; font-size:14px;
  padding:0 2px; }
.hcard .rm:hover { color:var(--ink); }
@media (max-width: 900px) { #ivx-app { display:block; height:auto; } #ivx-left { border-right:0; } }
"""

UI_JS = open("ui.js").read()

UI_JS = UI_JS.replace("VERSION", VER)
UI_HTML = UI_HTML.replace("VERSION", VER)

HEAD = f"<script>({UI_JS})();</script>"

with gr.Blocks(css=UI_CSS + "\n#ivx-bridge, #ivx-peekrow, #ivx-papersrow { display:none !important; }",
               head=HEAD, title=f"Dialectical abstracts {VER}") as demo:
    gr.HTML(UI_HTML)
    with gr.Row(elem_id="ivx-bridge"):
        bin_ = gr.Textbox(elem_id="ivx-bridge-in")
        bout = gr.Textbox(elem_id="ivx-bridge-out")
        bbtn = gr.Button("go", elem_id="ivx-bridge-btn")
    bbtn.click(generate, inputs=[bin_], outputs=[bout], api_name="generate")
    with gr.Row(elem_id="ivx-peekrow"):
        pj = gr.Textbox()
        po = gr.Textbox()
        pb = gr.Button("peek")
        pg = gr.Textbox()
    pb.click(peek, inputs=[pj], outputs=[po], api_name="peek")
    with gr.Row(elem_id="ivx-papersrow"):
        pp = gr.Textbox()
        ppb = gr.Button("papers")
    ppb.click(papers_payload, outputs=[pp], api_name="papers")
    gr.Button("ping", elem_id="ivx-ping", visible=False).click(ping, outputs=[pg], api_name="ping")

demo.queue().launch(ssr_mode=False)