"""MiniCPM-V: the local VLM — the find-and-point "eyes".

- ground_box: return the bounding box of a described object on a page image —
  the circle the agent (models/minicpm_agent, the text "brain") asks for.
- generate_answer: answers a question grounded in retrieved page images
  (one-shot; used by scripts/eval_answers_modal.py and the pipelines' .run()).
- describe_batch (+ describe_figure / describe_table wrappers): used only by
  the parsed ingest pipeline (on Modal) to turn Picture crops and Table
  markdown into searchable text, batched through chat()'s batched mode.

The model and tokenizer are module-level globals: ZeroGPU packs module-level
CUDA tensors at startup and shares them with the GPU worker, whereas function
arguments are pickled — and trust_remote_code model classes are not picklable.

All functions are plain (not @spaces.GPU): callers run them inside their own
GPU context — the ask pipelines' single @spaces.GPU call on the Space, or a
real CUDA process on Modal.
"""

from __future__ import annotations

import re

import torch
from PIL import Image
from transformers import AutoModel, AutoProcessor, AutoTokenizer

from core import tracing
from core.constants import (
    ANSWER_MAX_NEW_TOKENS,
    DESCRIBE_MAX_NEW_TOKENS,
    GROUND_BOX_MAX_NEW_TOKENS,
    GROUND_ENABLE_THINKING,
    GROUND_THINK_MAX_NEW_TOKENS,
    MINICPM_MODEL_ID,
    MINICPM_REVISION,
)
from core.vram import log_vram

PROMPT = (
    "You are a repair-manual assistant. The images are the manual pages most "
    "relevant to the user's question, each preceded by its label (manual name "
    "and page number).\n\n"
    "Answer the question using ONLY what is printed on these pages, following "
    "these rules:\n"
    "1. If the answer is a procedure, reproduce EVERY step in order, numbered "
    "exactly as in the manual. Never skip, merge, or summarize steps. Keep each "
    "step's notes, model exceptions, and specifications (e.g. torque values) "
    "with that step, exactly as printed.\n"
    "2. Quote exact values (torques, clearances, part numbers, capacities) as "
    "printed, including units.\n"
    "3. End with the page label(s) you used, e.g. (Manual — p.238). If a "
    "procedure clearly continues on a page you were not given, say so.\n"
    "4. If the pages do not contain the answer, say so instead of guessing.\n"
    "5. Start directly with the answer — no preamble like 'Based on the "
    "provided pages'.\n\n"
    "Question: {question}"
)

# Visual grounding for "circle the <thing>": the mechanic asks to circle
# something on the page they are viewing. MiniCPM-V grounding replies in
# <box>x1 y1 x2 y2</box> form with coordinates normalized to 0-1000.
# Abstention-first: the mechanic's circle target may not be a circleable thing on
# the page they're viewing (it's only a breadcrumb/heading, or simply absent), and
# a confident wrong box reads to them as "it pointed at the heading." So the FIRST
# job is to decide present-vs-not and say NOT FOUND otherwise — which also lets the
# 1B retry/relocate. Measured (no-think, eval_grounding_modal.py, n=30): vs the
# prior prompt this took NOT-FOUND recall 0.47->1.0 and breadcrumb-trap 0/5->5/5
# with no latency cost and without hurting same-page positives — i.e. it recovers
# thinking-mode's abstention without thinking-mode's token budget.
GROUND_PROMPT = (
    "The image is one page of a repair manual. A mechanic wants you to circle "
    "{query!r} on THIS page.\n"
    "FIRST decide whether {query!r} is actually on this page as a thing you can "
    "circle: a component drawn in a figure, a part with its label, a table row, "
    "or a specification value. The WORDS appearing is not enough — a section "
    "heading, a breadcrumb or navigation path (like 'A > B > C'), a page title, "
    "a caption, or a running header that merely NAMES {query!r} is not the thing "
    "itself. If the actual component/value/row is not on this page — even when "
    "its name shows up in a heading or breadcrumb — reply exactly: NOT FOUND.\n"
    "If it IS here, locate it precisely:\n"
    "- In an exploded or assembly diagram, parts carry callout numbers/letters on "
    "leader lines and a legend says what each number is. Find {query!r} in the "
    "legend to get its number, then follow that number's leader line to the part "
    "in the drawing and box THAT part — never the legend text or the number.\n"
    "- For a torque or specification, box the VALUE with its label/units (e.g. the "
    "number and N·m), not the whole table.\n"
    "- For a part whose name is printed as a label beside its own drawing, box the "
    "drawn part — not a heading or list elsewhere on the page.\n"
    "Box ONLY that one thing, as TIGHTLY as possible — just the part/value itself, "
    "never the whole figure, a group of parts, or the page. A box wider than about "
    "half the page is almost always wrong.\n"
    "Reply with ONLY the box, as <box>x1 y1 x2 y2</box>: four integers normalized "
    "0-1000 (x left→right, y top→bottom) over the whole page. If {query!r} is not "
    "a circleable thing on this page, reply exactly: NOT FOUND"
)

# Visual page rerank: score the search shortlist by LOOKING at the page images,
# instead of re-judging from page text (a 1B text rerank measured worse than raw
# ColEmbed top-1, because repair pages are figure-heavy and text retrieval
# surfaces index/spec pages that merely name-drop the part).
PAGE_RERANK_PROMPT = (
    "The image is one page of a repair manual. A mechanic wants: {query!r}\n"
    "How directly does THIS page give them what they need — the procedure, "
    "component, diagram, table, or value involved (judge by topic, not exact "
    "wording)? Reply with ONLY one digit 0-5: 5 = this is exactly the page, "
    "0 = unrelated."
)

FIGURE_PROMPT = (
    "The image is a figure cropped from a page of a repair manual. Context "
    "from the page it appears on:\n{context}\n\n"
    "Using the context's terminology, write a search-index description of the "
    "figure in 1-3 sentences: name the specific component or assembly shown "
    "(and which models it applies to, if stated), the kind of view or diagram, "
    "and what the labeled parts or callouts are. Be definite — never write "
    "'likely', 'possibly' or 'appears to'. Example of the expected style:\n"
    "'Dimensional reference drawing of the cross-type universal joint used on "
    "the Pn35-50 and Cu35-55 propeller shafts: two yokes joined by a spider "
    "and four bearing cups, shown in side and end views with the overall "
    "length, width, and height dimensions.'\n"
    "Start directly with the description — no preamble."
)

TABLE_PROMPT = (
    "Below is a table extracted from a repair manual as markdown, followed by "
    "context from the page it appears on.\n\nTable:\n{markdown}\n\n"
    "Context:\n{context}\n\n"
    "Using the context's terminology, write a search-index description of the "
    "table in 1-3 sentences: what the table is for, which components, models "
    "or operations it covers, and what values it lists (with units). Example "
    "of the expected style:\n"
    "'Propeller shaft specifications by forklift model (Pn35-50, Pn60-80, "
    "Cu35-55, Cu60/70): the universal joint type, the three shaft length "
    "dimensions in millimeters and inches, and whether the upper and lower "
    "propeller shaft covers are fitted.'\n"
    "Start directly with the description — no preamble."
)

_MODEL = (
    AutoModel.from_pretrained(
        MINICPM_MODEL_ID,
        revision=MINICPM_REVISION,
        trust_remote_code=True,
        dtype=torch.bfloat16,
        attn_implementation="sdpa",
    )
    .to("cuda")
    .eval()
)
_TOKENIZER = AutoTokenizer.from_pretrained(
    MINICPM_MODEL_ID, revision=MINICPM_REVISION, trust_remote_code=True
)
# Pre-build the processor chat() would otherwise lazily create per GPU worker
# (it caches on this exact attribute, see modeling_minicpmv.chat).
_MODEL.processor = AutoProcessor.from_pretrained(
    MINICPM_MODEL_ID, revision=MINICPM_REVISION, trust_remote_code=True
)
log_vram("load-minicpm-v")


def generate_answer(question: str, pages: list[tuple[str, Image.Image]]) -> str:
    """pages: [(label, page image)] in retrieval order. Must run on GPU
    (called from within a @spaces.GPU context)."""
    content = []
    for label, img in pages:  # chat() accepts interleaved strings and PIL images
        content.append(f"[{label}]")
        content.append(img.convert("RGB"))
    content.append(PROMPT.format(question=question))
    with torch.no_grad():
        answer = _MODEL.chat(
            msgs=[{"role": "user", "content": content}],
            tokenizer=_TOKENIZER,
            enable_thinking=False,
            max_new_tokens=ANSWER_MAX_NEW_TOKENS,
        )
    return str(answer).strip()


def ground_box(
    image: Image.Image,
    query: str,
    enable_thinking: bool | None = None,
    prompt: str | None = None,
) -> tuple[tuple[float, float, float, float] | None, str]:
    """(bbox, raw reply): the bounding box of the described object on a page
    image, in that image's pixel coordinates — or None when the model can't
    place it (no box in the reply, or a degenerate one). Must run on GPU.

    enable_thinking lets the model reason (legend → callout-number →
    leader-line → part) before committing to a box. None defers to the
    GROUND_ENABLE_THINKING default; the UI settings panel passes an explicit
    bool per request.

    prompt overrides GROUND_PROMPT (a template with a {query} field) — used by
    the grounding eval to A/B prompt variants without editing this module."""
    think = GROUND_ENABLE_THINKING if enable_thinking is None else enable_thinking
    template = prompt or GROUND_PROMPT
    # One `generation` (the VLM "eyes" placing the box): query in, raw reply out.
    with tracing.generation(
        "ground-circle",
        model=MINICPM_MODEL_ID,
        input=query,
        metadata={"thinking": bool(think)},
    ) as gen:
        with torch.no_grad():
            out = _MODEL.chat(
                msgs=[
                    {
                        "role": "user",
                        "content": [image.convert("RGB"), template.format(query=query)],
                    }
                ],
                tokenizer=_TOKENIZER,
                # The think trace needs room (a bare box fits in 64 tokens; a think
                # trace does not) or it gets cut off before emitting the box — so
                # the token budget tracks the flag.
                enable_thinking=think,
                max_new_tokens=(
                    GROUND_THINK_MAX_NEW_TOKENS if think else GROUND_BOX_MAX_NEW_TOKENS
                ),
            )
        raw = str(out).strip()
        if gen is not None:
            gen.update(output=raw)
    # The thinking trace can mention "not found" mid-reasoning ("at first this
    # looked not found, but…"), so test only the FINAL answer after </think>,
    # not the whole reply.
    final = raw.rsplit("</think>", 1)[-1]
    if "NOT FOUND" in final.upper():
        return None, raw
    # Read the coordinates from the <box>…</box> tag specifically. MiniCPM-V
    # may prefix a <ref>…</ref> (e.g. echoing "5. Rod"); that digit would
    # otherwise be grabbed as the first coordinate and shift the whole box.
    # Search `final` so digits inside the reasoning can't be mistaken for
    # coordinates.
    m = re.search(r"<box>(.*?)</box>", final, re.IGNORECASE | re.DOTALL)
    nums = re.findall(r"\d+(?:\.\d+)?", m.group(1) if m else final)
    if len(nums) < 4:
        return None, raw
    x1, y1, x2, y2 = (min(1000.0, max(0.0, float(n))) for n in nums[:4])
    if x2 - x1 < 1 or y2 - y1 < 1:
        return None, raw
    w, h = image.size
    return (x1 * w / 1000, y1 * h / 1000, x2 * w / 1000, y2 * h / 1000), raw


def rerank_pages(images: list[Image.Image], query: str) -> tuple[int, list[int]]:
    """Pick the best of the search shortlist by LOOKING at the page images.
    Each candidate is scored 0-5 for how directly it answers the query, in ONE
    batched chat() call; the highest score wins, ties broken by retrieval order
    (so on a tie it never does worse than ColEmbed). Returns (index into images,
    the per-page scores; -1 where the reply had no digit). Must run on GPU."""
    prompt = PAGE_RERANK_PROMPT.format(query=query)
    msgs = [
        [{"role": "user", "content": [img.convert("RGB"), prompt]}] for img in images
    ]
    with torch.no_grad():
        out = _MODEL.chat(
            msgs=msgs,
            tokenizer=_TOKENIZER,
            enable_thinking=False,
            max_new_tokens=4,
        )
    scores = []
    for a in out:
        m = re.search(r"\d", str(a))
        scores.append(int(m.group()) if m else -1)
    # max() returns the FIRST argmax, so equal scores keep ColEmbed's order.
    best = max(range(len(scores)), key=lambda i: scores[i]) if scores else 0
    return best, scores


def describe_batch(jobs: list[tuple[str, object, str]]) -> list[str]:
    """Searchable descriptions for a batch of parsed elements in ONE chat()
    call (its batched mode: msgs = list of conversations). Each job is
    ("figure", crop image, context) or ("table", markdown, context); the
    context (manual name, section heading, caption, page text) is what lets
    MiniCPM use the manual's own terminology. Returns descriptions in job
    order. Parsed ingest only; must run on GPU."""
    msgs = []
    for kind, payload, context in jobs:
        if kind == "figure":
            content = [payload.convert("RGB"), FIGURE_PROMPT.format(context=context)]
        else:
            content = [TABLE_PROMPT.format(markdown=payload, context=context)]
        msgs.append([{"role": "user", "content": content}])
    with torch.no_grad():
        out = _MODEL.chat(
            msgs=msgs,
            tokenizer=_TOKENIZER,
            enable_thinking=False,
            max_new_tokens=DESCRIBE_MAX_NEW_TOKENS,
        )
    return [str(answer).strip() for answer in out]


def describe_figure(image: Image.Image, context: str) -> str:
    """Single-element convenience wrapper around describe_batch."""
    return describe_batch([("figure", image, context)])[0]


def describe_table(markdown: str, context: str) -> str:
    """Single-element convenience wrapper around describe_batch."""
    return describe_batch([("table", markdown, context)])[0]