import spaces  # must be imported before torch / transformers

import html
import time

import gradio as gr
import torch
from huggingface_hub import snapshot_download
from transformers import AutoProcessor, Qwen3VLForConditionalGeneration

from docision._runtime.hf import HFDocisionEngine
from docision.engine import DocisionEngine
from docision.io import DEFAULT_CONTEXT

MODEL_ID = "SoMarkAI/Docision-1.0"
# Reference defaults from the official CLI (docision/cli.py)
MIN_PIXELS = 3136
MAX_PIXELS = 501760
MAX_INPUT_TOKENS = 4000

model_path = snapshot_download(MODEL_ID, allow_patterns=["*.json", "*.jinja", "*.safetensors"])
processor = AutoProcessor.from_pretrained(model_path)
model, loading_info = Qwen3VLForConditionalGeneration.from_pretrained(
    model_path, dtype=torch.bfloat16, attn_implementation="sdpa", output_loading_info=True
)
# Same strictness as the official runtime: refuse incomplete checkpoint loads.
if any(loading_info.get(k) for k in ("missing_keys", "unexpected_keys", "mismatched_keys", "error_msgs")):
    raise RuntimeError(f"Incomplete Docision checkpoint load: {loading_info}")
print("[docision] checkpoint loaded cleanly")
model = model.to("cuda")
model.eval()

# Build the official Docision engine around the already-loaded model
# (DocisionEngine.__init__ would load with device_map="auto", which bypasses ZeroGPU).
engine = DocisionEngine.__new__(DocisionEngine)
HFDocisionEngine.__init__(
    engine, model_path, model=model, processor=processor,
    numerics="native", max_input_tokens=MAX_INPUT_TOKENS,
)
engine.builder.processor.image_processor.size = {
    "shortest_edge": MIN_PIXELS,
    "longest_edge": MAX_PIXELS,
}


def parse_questions(text: str) -> list[dict]:
    """Parse blocks separated by blank lines: first line = question, other lines = candidates."""
    blocks, current = [], []
    for raw in (text or "").splitlines():
        line = raw.strip()
        if not line:
            if current:
                blocks.append(current)
                current = []
            continue
        current.append(line)
    if current:
        blocks.append(current)
    if not blocks:
        raise gr.Error("Add at least one question followed by two or more candidate lines.")
    questions = []
    for i, block in enumerate(blocks, 1):
        question = block[0]
        for prefix in ("Q:", "Q："):
            if question.startswith(prefix):
                question = question[len(prefix):].strip()
        candidates = []
        for line in block[1:]:
            if line[:2] in ("- ", "* ", "• "):
                line = line[2:].strip()
            if line:
                candidates.append(line)
        if len(candidates) < 2:
            raise gr.Error(f"Question {i} (“{question}”) needs at least two candidate lines.")
        if len(set(candidates)) != len(candidates):
            raise gr.Error(f"Question {i} has duplicate candidates.")
        questions.append({"id": f"q{i}", "question": question, "candidates": candidates})
    return questions


def render_html(questions: list[dict], results: list[dict], seconds: float) -> str:
    parts = ['<div class="dc-wrap">']
    for q, r in zip(questions, results):
        decision = r["decision"]
        parts.append('<div class="dc-card">')
        parts.append(f'<div class="dc-q">{html.escape(q["question"])}</div>')
        parts.append(
            f'<div class="dc-dec">Decision: <b>[{html.escape(decision["label"])}] '
            f'{html.escape(decision["text"])}</b> &nbsp;·&nbsp; {decision["probability"]*100:.1f}%</div>'
        )
        for c in r["candidates"]:
            pct = c["probability"] * 100
            win = " dc-win" if c["label"] == decision["label"] else ""
            parts.append(
                f'<div class="dc-row{win}">'
                f'<div class="dc-label">[{html.escape(c["label"])}] {html.escape(c["text"])}</div>'
                f'<div class="dc-barbox"><div class="dc-bar" style="width:{max(pct, 0.3):.2f}%"></div></div>'
                f'<div class="dc-pct">{pct:.2f}%</div></div>'
            )
        parts.append("</div>")
    parts.append(f'<div class="dc-foot">{len(results)} question(s) scored in {seconds:.2f}s on GPU</div></div>')
    return "".join(parts)


def _gpu_duration(image, questions_text, *args, **kwargs):
    # Measured ~0.5-0.7s per question (incl. cold start); keep a small fixed margin.
    n = sum(1 for block in (questions_text or "").split("\n\n") if block.strip())
    return min(90, 10 + 2 * max(n, 1))


@spaces.GPU(duration=_gpu_duration)
def decide(
    image: str,
    questions_text: str,
    context: str = DEFAULT_CONTEXT,
    mode: str = "causal",
) -> tuple[str, list]:
    """Score candidate answers to one or more questions about a document image with Docision-1.0.

    Args:
        image: path to the document image (PNG/JPEG/WebP).
        questions_text: question blocks separated by blank lines; the first line of each block is the
            question, every following line is one candidate answer.
        context: instruction prefix shared by all questions.
        mode: "causal" (default; candidates see each other) or "isolated" (each candidate attends only to the shared prefix).

    Returns:
        An HTML visualization and the raw per-question results (candidate probabilities + decision).
    """
    if not image:
        raise gr.Error("Upload a document image first.")
    questions = parse_questions(questions_text)
    document = {
        "image": image,
        "context": context if context and context.strip() else DEFAULT_CONTEXT,
        "questions": questions,
    }
    t0 = time.perf_counter()
    try:
        results = engine.score_document(document, mode=mode, prefix_cache=False, batch_size=1)
    except ValueError as e:
        raise gr.Error(str(e))
    seconds = time.perf_counter() - t0
    print(f"[docision] scored {len(questions)} question(s) in {seconds:.2f}s")
    raw = [
        {
            "id": q["id"],
            "question": q["question"],
            "decision": {k: r["decision"][k] for k in ("label", "text", "probability", "index")},
            "candidates": r["candidates"],
            "prompt_tokens": r["prompt_tokens"],
        }
        for q, r in zip(questions, results)
    ]
    return render_html(questions, results, seconds), raw


EX_EN = """What is the RR for severely malnourished in the table?
3.4
4.8
2.1

What is used as the reference in the Stage at Diagnosis row?
Stages I and II
Newly diagnosed
Well-nourished

Are the 95% CIs for moderately malnourished and stage at diagnosis the same?
Different
The same

Which items in the table have P-values marked “< 0.001”?
Moderately malnourished and severely malnourished
Severely malnourished and treatment history
Stage at diagnosis and treatment history
Moderately malnourished and stage at diagnosis"""

EX_ZH = """表中重度营养不良（Severely malnourished）的RR是多少？
3.4
4.8
2.1

诊断时分期（Stage at Diagnosis）这一行采用什么作为参照？
I期和II期
新诊断
营养良好

中度营养不良与诊断时分期的95% CI是否相同？
不同
相同"""

EX_CLS = """What type of document is this?
A statistical results table from a research paper
An invoice
A bank statement
A handwritten letter

Does the table report relative risks with confidence intervals?
Yes
No"""

CSS = """
#col-container { max-width: 1200px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
.dc-wrap { display: flex; flex-direction: column; gap: 12px; }
.dc-card { border: 1px solid var(--border-color-primary); border-radius: 10px; padding: 12px 14px; background: var(--block-background-fill); }
.dc-q { font-weight: 600; margin-bottom: 4px; color: var(--body-text-color); }
.dc-dec { font-size: 0.9em; margin-bottom: 8px; color: var(--body-text-color-subdued); }
.dc-row { display: grid; grid-template-columns: minmax(140px, 2fr) 3fr 70px; gap: 8px; align-items: center; padding: 3px 0; font-size: 0.92em; color: var(--body-text-color); }
.dc-barbox { background: var(--background-fill-secondary); border-radius: 6px; height: 14px; overflow: hidden; }
.dc-bar { height: 100%; background: var(--neutral-400); border-radius: 6px; }
.dc-win .dc-bar { background: var(--color-accent); }
.dc-win .dc-label { font-weight: 600; }
.dc-pct { text-align: right; font-variant-numeric: tabular-nums; }
.dc-foot { font-size: 0.8em; color: var(--body-text-color-subdued); }
"""

with gr.Blocks(title="Docision-1.0") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            "# 📑 Docision-1.0\n"
            "Document decisions with candidate scoring: upload a document image, write one or more questions, "
            "each with its own candidate answers, and get the chosen candidate plus a probability for every option. "
            "Fine-tuned from Qwen3-VL-4B-Instruct.  \n"
            "[Model](https://huggingface.co/SoMarkAI/Docision-1.0) · "
            "[GitHub](https://github.com/SoMarkAI/docision) · "
            "[Benchmark](https://huggingface.co/datasets/SoMarkAI/doc-decision-bench)"
        )
        with gr.Row():
            with gr.Column(scale=1):
                image = gr.Image(label="Document image", type="filepath", format="png", height=320)
                questions_text = gr.Textbox(
                    label="Questions & candidates",
                    info="One block per question, separated by a blank line. First line = question, "
                         "each following line = one candidate (2–128 per question).",
                    lines=12,
                    placeholder="Which category does this document belong to?\nInvoice\nReceipt\nContract\n\n"
                                "Is the document signed?\nYes\nNo",
                )
                run = gr.Button("Decide", variant="primary")
                with gr.Accordion("Advanced options", open=False):
                    context = gr.Textbox(label="Context / instructions (shared by all questions)",
                                         value=DEFAULT_CONTEXT, lines=2)
                    mode = gr.Radio(["causal", "isolated"], value="causal", label="Attention mode",
                                    info="isolated: each candidate only sees the shared prompt, not the other candidates.")
            with gr.Column(scale=1):
                result_html = gr.HTML(label="Decisions")
                result_json = gr.JSON(label="Raw output")

        gr.Examples(
            examples=[
                ["examples/validation-table.jpg", EX_EN],
                ["examples/validation-table.jpg", EX_ZH],
                ["examples/validation-table.jpg", EX_CLS],
            ],
            inputs=[image, questions_text],
            outputs=[result_html, result_json],
            fn=decide,
            cache_examples=True,
            cache_mode="lazy",
        )
        gr.Markdown(
            "<small>Example image: Table 3 from Gupta et al., *Can subjective global assessment of nutritional "
            "status predict survival in ovarian cancer?*, Journal of Ovarian Research 1, 5 (2008), "
            "[CC BY 2.0](https://creativecommons.org/licenses/by/2.0/). Probabilities are relative scores over the "
            "candidates, not calibrated confidences.</small>"
        )

    run.click(decide, inputs=[image, questions_text, context, mode],
              outputs=[result_html, result_json], api_name="decide")

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)