"""RAP demo UI: PDF -> assets -> model HTML -> rendered poster.  Pipeline unchanged; UI redesigned 2026-10-15."""
import json
import re
import tempfile
import time
from pathlib import Path
import gradio as gr
import torch
import spaces
from pdf_pipeline import prepare, chat, CANVASES, DEFAULT_INSTRUCTION
from paper_sections import budget_body
from poster_render import render

ROOT = Path(__file__).resolve().parent
ARTIFACTS = Path(tempfile.mkdtemp(prefix='rap_demo_'))
WINDOW = 16384
MAX_STEPS = 10   # chained ZeroGPU steps per run
GPU_LADDER = (600, 300, 120, 60)   # requested step length; ZeroGPU refuses what the visitor's quota cannot cover
QUOTA_HINT = (' — this 27B BF16 model needs the full 96 GB ZeroGPU (2x quota) and a poster takes roughly 15–25 min of quota. '
              'Daily ZeroGPU quota: PRO 40 min, free account 5 min, anonymous 2 min, so sign in with a PRO account to finish a poster.')

CSS = '''
#hero .links { margin: 14px 0 0; font-size: 14px; opacity: .92 } #hero .links a { color: #fff; text-decoration: underline; text-underline-offset: 3px }
.gradio-container {max-width: 1360px !important; margin: 0 auto;}
#hero {padding: 26px 30px 22px; border-radius: 16px; margin-bottom: 6px;
       background: linear-gradient(120deg, #1e1b4b 0%, #4c3d99 60%, #6d5bd0 100%); color: #fff;}
#hero h1 {margin: 0 0 6px; font-size: 30px; letter-spacing: .2px; color: #fff;}
#hero p {margin: 0; color: #e6e2ff; font-size: 15px;}
#hero a {color: #fff; text-decoration: underline;}
.badge {display: inline-block; margin: 10px 8px 0 0; padding: 3px 10px; border-radius: 999px;
        background: rgba(255,255,255,.14); font-size: 12px; color: #fff;}
#status {min-height: 24px; font-size: 14px;}
#poster img {border-radius: 10px; box-shadow: 0 10px 30px rgba(20,16,60,.12);}
#examples .thumbnail-item {border-radius: 10px;}
footer {display: none !important;}
'''

def load_examples():
    """examples/<name>/: meta.json, paper.pdf, poster.png, poster.html, prompt.json (the exact benchmark prompt),
    thumbnails/ (what the model sees) and assets/ (full-resolution crops for rendering)."""
    import hashlib
    out = []
    for d in sorted((ROOT / 'examples').glob('*/')):
        meta = d / 'meta.json'
        if not meta.exists() or not (d / 'poster.png').exists() or not (d / 'paper.pdf').exists():
            continue
        m = json.loads(meta.read_text(encoding='utf-8'))
        out.append({'dir': d, 'title': m.get('title', d.name), 'short': m.get('short', m.get('title', d.name))[:40], 'venue': m.get('venue', ''),
                    'paper_id': m.get('paper_id', d.name), 'instruction': m.get('instruction', DEFAULT_INSTRUCTION), 'spec': m.get('spec', ''),
                    'canvas': m.get('canvas', 'Landscape · 3:2'), 'columns': m.get('columns', 3), 'accent': m.get('accent', '#584287'),
                    'prompt': d / 'prompt.json' if (d / 'prompt.json').exists() else None,
                    'pdf_sha': hashlib.sha256((d / 'paper.pdf').read_bytes()).hexdigest()})
    return out

def fileout(path):
    return gr.update(value=str(path), visible=True) if path else gr.update(value=None, visible=False)

def build(model, processor):
    def match_example(pdf, instruction, custom_spec):
        """The exact benchmark prompt is used when the PDF is one of the examples and instruction + spec are unchanged."""
        import hashlib
        try:
            sha = hashlib.sha256(Path(pdf).read_bytes()).hexdigest()
        except OSError:
            return None
        for ex in examples:
            if ex['prompt'] and ex['pdf_sha'] == sha and instruction.strip() == ex['instruction'].strip() and custom_spec.strip() == ex['spec'].strip():
                return ex
        return None

    def analyze_example(ex):
        """Reproduce a benchmark case: the stored prompt (paper text, asset list, thumbnails, instruction, layout spec)."""
        d = ex['dir']; pr = json.loads((d / 'prompt.json').read_text(encoding='utf-8'))
        system, user = pr['messages'][0]['content'], pr['messages'][1]['content']
        parts = user.split('<image>'); images = [str(d / im) for im in pr['images']]
        content = []
        for k, text in enumerate(parts):
            if text: content.append({'type': 'text', 'text': text})
            if k < len(images): content.append({'type': 'image', 'image': images[k]})
        messages = [{'role': 'system', 'content': [{'type': 'text', 'text': system}]}, {'role': 'user', 'content': content}]
        assets = []
        for im in pr['images']:
            aid = Path(im).stem; f = d / 'assets' / f'{aid}.png'
            kind = aid.rsplit('_', 1)[0]
            assets.append({'id': aid, 'kind': kind, 'file': str(f if f.exists() else d / im), 'thumb': str(d / im), 'page': None})
        w, h = (int(v) for v in re.match(r'(\d+)x(\d+)', ex['canvas']).groups()) if re.match(r'\d+x\d+', ex['canvas']) else CANVASES.get(ex['canvas'], (2400, 1600))
        job = Path(tempfile.mkdtemp(prefix='example_', dir=ARTIFACTS))
        return {'title': ex['title'], 'authors': '', 'assets': assets, 'instruction': ex['instruction'], 'spec': ex['spec'], 'canvas': [w, h],
                'messages_pre': messages, 'job': str(job), 'example': ex['paper_id'], 'sections': [], 'body': ''}

    def analyze(pdf, instruction, canvas, accent, columns, custom_spec):
        ex = match_example(pdf, instruction, custom_spec)
        if ex is not None:
            p = analyze_example(ex)
            inputs = processor.apply_chat_template(chat(p), tokenize=True, add_generation_prompt=True, enable_thinking=False, return_dict=True, return_tensors='pt')
            p['input_tokens'] = int(inputs['input_ids'].shape[1]); p['continuation_ids'] = []; p['raw'] = ''; p['done'] = False
            return p
        job = Path(tempfile.mkdtemp(prefix='paper_', dir=ARTIFACTS))
        p = prepare(pdf, job, instruction, canvas, accent, columns, custom_spec)
        budget = 22000
        while True:
            inputs = processor.apply_chat_template(chat(p), tokenize=True, add_generation_prompt=True,
                                                   enable_thinking=False, return_dict=True, return_tensors='pt')
            n = inputs['input_ids'].shape[1]
            if n <= 10000:
                break
            if budget <= 3500:
                raise ValueError('The paper and its figures exceed the demo input budget.')
            budget = max(3500, int(budget * .8))
            p['body'], _ = budget_body(p['sections'], total=budget, abstract_missing=True)
        p['input_tokens'] = n; p['continuation_ids'] = []; p['raw'] = ''; p['done'] = False
        return p

    def gpu_seconds(p):
        return int(p.get('gpu_seconds', 120))

    @spaces.GPU(size='xlarge', duration=gpu_seconds)   # 27B BF16 needs the full 96 GB GPU (2x quota)
    def gpu_chunk(p):
        budget = max(30, gpu_seconds(p) - 90)   # leave room for moving the weights to the GPU and the prefill
        inputs = processor.apply_chat_template(chat(p), tokenize=True, add_generation_prompt=True,
                                               enable_thinking=False, return_dict=True, return_tensors='pt')
        prefix = p.get('continuation_ids', [])
        if prefix:
            inputs['input_ids'] = torch.cat([inputs['input_ids'], torch.tensor([prefix], dtype=torch.long)], dim=1)
            inputs['attention_mask'] = torch.ones_like(inputs['input_ids'])
        length = int(inputs['input_ids'].shape[1])
        remaining = WINDOW - length
        if remaining <= 0:
            raise ValueError('The 16K context budget has been exhausted.')
        inputs = inputs.to('cuda')
        start = time.monotonic()
        with torch.inference_mode():
            output = model.generate(**inputs, max_new_tokens=min(8192, remaining), max_time=budget, do_sample=False,
                                    use_cache=True, temperature=None, top_p=None, top_k=None, repetition_penalty=1.1)
        new_ids = output[0, length:].detach().cpu().tolist()
        all_ids = prefix + new_ids
        text = processor.decode(all_ids, skip_special_tokens=True)
        eos = model.generation_config.eos_token_id
        eos = eos if isinstance(eos, list) else [eos]
        complete = bool(re.search(r'</html\s*>', text, re.I))
        stopped = complete or bool(new_ids and new_ids[-1] in eos) or length + len(new_ids) >= WINDOW
        result = {'ids': all_ids, 'raw': text, 'done': stopped, 'complete': complete, 'seconds': time.monotonic() - start,
                  'new_tokens': len(new_ids), 'total_tokens': p['input_tokens'] + len(all_ids)}
        del inputs, output
        return result

    HIDE = gr.update(value=None, visible=False)

    def gallery_of(p):
        return [(a['thumb'], f"{a['id']} · page {a['page']}" if a.get('page') else a['id']) for a in p['assets']]

    def analyze_evt(pdf, instruction, canvas, accent, columns, custom_spec, progress=gr.Progress()):
        """Request 1 of a run: parse the PDF on CPU. GPU steps run in their own requests (see step)."""
        if not pdf:
            raise gr.Error('Upload a paper PDF first.')
        progress(.1, desc='Reading the paper (LayoutLMv3 on CPU, up to a few minutes)…')
        try:
            p = analyze(pdf, instruction, canvas, accent, columns, custom_spec)
        except (ValueError, RuntimeError) as e:
            raise gr.Error(str(e)) from None
        p = dict(p, step=0, paused=False)
        head = (f"Benchmark case **{p['title'][:80]}** · exact test prompt, {len(p['assets'])} figures/tables" if p.get('example')
                else f"Parsed **{p['title'][:90]}** · {len(p['assets'])} figures/tables found")
        return (p, f"{head} · generating (step 1)…",
                None, HIDE, HIDE, gallery_of(p), p['spec'], '', gr.update(visible=False))

    def step(p, progress=gr.Progress()):
        """One ZeroGPU step per request: the ZeroGPU proxy token is tied to the request and expires after a few minutes,
        so a single long request (parse + several 120 s steps) fails with 'Expired ZeroGPU proxy token'."""
        if not p or p.get('done') or p.get('paused'):
            return tuple(gr.skip() for _ in range(9))
        k = p.get('step', 0) + 1
        gallery = gallery_of(p)
        if k > MAX_STEPS:
            return (dict(p, paused=True), f"Stopped after {MAX_STEPS} steps without a complete poster.", None, HIDE, HIDE,
                    gallery, p['spec'], p['raw'], gr.update(visible=True))
        progress(min(.1 + .12 * k, .9), desc=f'Generating poster HTML (step {k})…')
        r = None
        for secs in [d for d in GPU_LADDER if d <= p.get('gpu_seconds', GPU_LADDER[0])]:
            p = dict(p, gpu_seconds=secs)
            try:
                r = gpu_chunk(p)
                break
            except Exception as e:                                         # noqa: BLE001  (quota, duration limit, timeout, …)
                msg = str(e)
                if 'larger than the maximum allowed' in msg and secs != GPU_LADDER[-1]:
                    continue                                               # ask for a shorter GPU slot
                hint = QUOTA_HINT if ('maximum allowed' in msg or 'quota' in msg.lower()) else ''
                return (dict(p, paused=True), f"Paused after step {k - 1}: {msg[:200]}{hint}", None, HIDE, HIDE,
                        gallery, p['spec'], p['raw'], gr.update(visible=bool(p.get('continuation_ids'))))
        p = dict(p, continuation_ids=r['ids'], raw=r['raw'], done=r['done'], step=k)
        tokens = f"{len(r['ids']):,} tokens · {r['new_tokens'] / max(r['seconds'], .01):.0f} tok/s · {p['gpu_seconds']} s GPU step"
        if not (r['complete'] or p['done']):
            return (p, f"Step {k} done · {tokens} · continuing (step {k + 1})…", None, HIDE, HIDE,
                    gallery, p['spec'], p['raw'], gr.update(visible=False))
        progress(.95, desc='Rendering…')
        try:
            png, html, archive, meta = render(p['raw'], p)
        except Exception as e:                                             # noqa: BLE001
            return (dict(p, paused=True), f"The model stopped before a complete poster ({str(e)[:160]}).", None, HIDE, HIDE,
                    gallery, p['spec'], p['raw'], gr.update(visible=not p.get('done')))
        note = '' if meta['complete_html'] else ' · output incomplete'
        return (dict(p, done=True), f"Done · {tokens}{note}", png, fileout(html), fileout(archive),
                gallery, p['spec'], p['raw'], gr.update(visible=False))

    def resume(p):
        if not p or p.get('done'):
            raise gr.Error('Nothing to continue. Click Generate poster.')
        return dict(p, paused=False)

    examples = load_examples()
    def pick_example(evt: gr.SelectData):
        ex = examples[evt.index]; d = ex['dir']
        html = d / 'poster.html'; archive = d / 'poster_bundle.zip'
        return (str(d / 'paper.pdf'), ex['instruction'], ex['canvas'], ex['accent'], ex['columns'], ex['spec'],
                f"Benchmark case · {ex['title']}" + (f" ({ex['venue']})" if ex['venue'] else '') + ' · click **Generate poster** to reproduce it',
                str(d / 'poster.png'), fileout(html if html.exists() else None), fileout(archive if archive.exists() else None))

    with gr.Blocks(title='RAP · Paper to Poster', css=CSS, delete_cache=(3600, 86400),
                   theme=gr.themes.Soft(primary_hue='violet', neutral_hue='slate', radius_size='lg')) as demo:
        gr.HTML('<div id="hero"><h1>RAP · Paper to Poster</h1>'
                '<p>Upload a research paper and get an editable academic poster: a 27B vision-language model '
                'reads the paper, picks figures and tables, and writes the poster as standalone HTML.</p>'
                '<span class="badge">Qwen3.8-27B · SFT + GRPO</span><span class="badge">BF16 · ZeroGPU</span>'
                '<span class="badge">Output: HTML + PNG</span>'
                '<p class="links"><a href="https://arxiv.org/abs/2610.23230" target="_blank">Paper</a> · '
                '<a href="https://reconstruct-a-poster.github.io/" target="_blank">Project page</a> · '
                '<a href="https://github.com/VincentHancoder/RAP" target="_blank">Code</a> · '
                '<a href="https://huggingface.co/VincentHancoder/RAP-Qwen3.8-27B-RL" target="_blank">Model</a> · '
                '<a href="https://huggingface.co/datasets/VincentHancoder/RAP" target="_blank">Dataset RAP-25K</a></p></div>')
        with gr.Row(equal_height=False):
            with gr.Column(scale=5, min_width=360):
                pdf = gr.File(label='Paper PDF', file_types=['.pdf'], type='filepath', height=110)
                instruction = gr.Textbox(label='Instruction', value=DEFAULT_INSTRUCTION, lines=3)
                with gr.Row():
                    canvas = gr.Dropdown(list(CANVASES), value='Landscape · 3:2', label='Canvas (preset or WxH)', allow_custom_value=True, scale=2)
                    columns = gr.Slider(1, 4, value=3, step=1, label='Columns', scale=2)
                    accent = gr.ColorPicker(value='#584287', label='Accent', scale=1)
                with gr.Accordion('Custom layout specification (optional)', open=False):
                    custom_spec = gr.Textbox(show_label=False, placeholder='Leave empty to let the demo plan the layout. Benchmark cases fill in their exact layout specification here.', lines=4)
                run_btn = gr.Button('Generate poster', variant='primary', size='lg')
                cont_btn = gr.Button('Continue generation', variant='secondary', visible=False)
                gr.Markdown('<small>Generation runs in ZeroGPU steps of up to 120 s; a full poster usually takes 3–6 steps. '
                            'Sign in to Hugging Face to use your GPU quota.</small>')
            with gr.Column(scale=7):
                status = gr.Markdown('Upload a PDF or pick an example below.', elem_id='status')
                poster = gr.Image(label='Poster', type='filepath', interactive=False, height=560, elem_id='poster', show_download_button=True)
                with gr.Row():
                    html_download = gr.File(label='Poster HTML', interactive=False, visible=False)
                    bundle = gr.File(label='HTML + figures (ZIP)', interactive=False, visible=False)
        if examples:
            gr.Markdown('### Examples')
            ex_gallery = gr.Gallery(value=[(str(e['dir'] / 'poster.png'), e['short']) for e in examples], columns=5, rows=1,
                                    height=230, object_fit='contain', show_label=False, elem_id='examples', allow_preview=False)
        with gr.Accordion('Details: extracted figures, layout plan, raw HTML', open=False):
            gallery = gr.Gallery(label='Figures and tables offered to the model', columns=6, height=220)
            spec_view = gr.Textbox(label='Layout specification', lines=4, interactive=False)
            code = gr.Code(language='html', label='Model output', interactive=False)
        with gr.Accordion('Notes', open=False):
            gr.Markdown('Research demo. Scanned PDFs are not supported; uploads are processed on this Space and are temporary. '
                        'Check every statement, number and caption before using a poster. '
                        '[Source and third-party notices](https://huggingface.co/spaces/VincentHancoder/RAP/tree/main).')
        state = gr.State()
        outs = [state, status, poster, html_download, bundle, gallery, spec_view, code, cont_btn]
        # Chain: one request for parsing, then one request per GPU step (fresh ZeroGPU token each time).
        # API clients: call /analyze once, then /step until the status starts with "Done", "Paused" or "Stopped".
        ev = run_btn.click(analyze_evt, [pdf, instruction, canvas, accent, columns, custom_spec], outs, api_name='analyze',
                           concurrency_limit=1)
        for k in range(MAX_STEPS):
            ev = ev.then(step, [state], outs, api_name='step' if k == 0 else False, concurrency_limit=1)
        ev = cont_btn.click(resume, [state], [state], api_name=False)
        for _ in range(MAX_STEPS):
            ev = ev.then(step, [state], outs, api_name=False, concurrency_limit=1)
        if examples:
            ex_gallery.select(pick_example, None, [pdf, instruction, canvas, accent, columns, custom_spec, status, poster, html_download, bundle], api_name=False)
    return demo, ARTIFACTS