pipelines/mock_ask.py

"""Mock store + find-and-point pipeline for local UI iteration (MOCK_MODELS=1).

Drop any PDF into MOCK_PDF_DIR (default app/data/mock_pdfs/) and the whole
find-and-point UX runs over real, rendered pages of that PDF — no GPU, no
model downloads, no HF library sync. The same MockStore instance backs both
approaches, so the manual dropdown, the page viewer, navigation, the three
router branches (go_to_section / point_here / search→classify→circle), the
circle overlay and the candidate-pages strip all exercise the real wiring;
only the router's tool choice, the page/bbox picks and the page classification
are faked (by simple keyword heuristics, so the branches are predictable in
tests).

Nothing here imports torch / spaces / the model modules (those load CUDA at
import), so this module is safe to load on a laptop.
"""

from __future__ import annotations

import glob
import hashlib
import os
import re
import time

from core.constants import MOCK_PDF_DIR
from core.pdf import page_count, page_size, render_pages
from core.sections import match_section
from core.store import slugify

# Canned section titles spread evenly over each PDF's pages, realistic enough
# to exercise "go to <section>" fuzzy matching in local tests.
MOCK_SECTION_TITLES = [
    "General Information",
    "Engine Mechanical System",
    "Engine Electrical System",
    "Fuel System",
    "Cooling System",
    "Transmission System",
    "Brake System — Bleeding and Adjustment",
    "Steering System",
    "Suspension System",
    "Body Electrical System",
]

MOCK_ANSWER = (
    "🧪 **Mock mode** — canned answer for UI iteration, not a real model "
    "response.\n\n"
    "1. Loosen the three retaining bolts in a star pattern (reassembly torque: "
    "24 N·m / 18 ft-lb).\n"
    "2. Withdraw the assembly and inspect the seal lip for wear or scoring.\n"
    "3. Refit in reverse order, checking clearances against the spec table.\n\n"
    "The pages this answer *would* be grounded in are shown on the right.\n\n"
    "({label})"
)


class MockStore:
    """A folder of PDFs exposed through the slice of the DocStore API that
    app.py and MockAskPipeline use. doc_id is the slug of the file name."""

    def __init__(self, pdf_dir: str = MOCK_PDF_DIR):
        self.pdf_dir = os.path.abspath(pdf_dir)
        os.makedirs(self.pdf_dir, exist_ok=True)

    def _docs(self) -> dict[str, dict]:
        """Re-scanned on every call so PDFs dropped in while the app is running
        appear after a 🔄 Sync library click (which just rebuilds the dropdown)."""
        docs: dict[str, dict] = {}
        for path in sorted(glob.glob(os.path.join(self.pdf_dir, "*.pdf"))):
            name = os.path.splitext(os.path.basename(path))[0]
            docs[slugify(name)] = {"name": name, "path": path}
        return docs

    def list_docs(self) -> list[dict]:
        return [
            {
                "doc_id": doc_id,
                "name": info["name"],
                "pages": page_count(info["path"]),
                "size_mb": os.path.getsize(info["path"]) / 1e6,
            }
            for doc_id, info in self._docs().items()
        ]

    def exists(self, doc_id: str) -> bool:
        return doc_id in self._docs()

    def pdf_path(self, doc_id: str) -> str | None:
        info = self._docs().get(doc_id)
        return info["path"] if info else None

    def sections(self, doc_id: str) -> list[dict]:
        """Canned sections spread evenly over the PDF's real pages — same
        shape as core.sections.sections_from_chunks."""
        info = self._docs().get(doc_id)
        if not info:
            return []
        n = page_count(info["path"])
        k = min(len(MOCK_SECTION_TITLES), n)
        starts = [round(i * n / k) + 1 for i in range(k)]
        return [
            {
                "title": MOCK_SECTION_TITLES[i],
                "page_start": starts[i],
                "page_end": (starts[i + 1] - 1) if i + 1 < k else n,
            }
            for i in range(k)
        ]



class MockAskPipeline:
    """Stateless mock matching the real ask pipelines' contracts: picks a
    deterministic spread of real pages and returns a canned answer."""

    def run(self, store: MockStore, question: str, doc_ids: list[str] | None, top_k: int):
        """Return (answer markdown, gallery [(image, caption)], page_refs
        [(doc_id, page)]) — same shape as the real ask pipelines."""
        question = (question or "").strip()
        if not question:
            raise ValueError("Please enter a question.")
        # Real questions take seconds on the GPU; the mock is instant, so the
        # loading indicator never shows. MOCK_DELAY (seconds) fakes that latency
        # for local UI work — e.g. MOCK_DELAY=2.
        time.sleep(float(os.environ.get("MOCK_DELAY", "0")))
        doc_id, info = self._pick_doc(store, doc_ids)
        pages = self._pick_pages(question, info["pages"], int(top_k))
        images = render_pages(store.pdf_path(doc_id), pages)
        labels = [f"{info['name']} — p.{p}" for p in pages]
        answer = MOCK_ANSWER.format(label=labels[0])
        gallery = [(img, f"{label} (mock)") for label, img in zip(labels, images)]
        page_refs = [(doc_id, p) for p in pages]
        return answer, gallery, page_refs

    def run_find(
        self,
        store: MockStore,
        request: str,
        doc_ids: list[str] | None,
        top_k: int,
        sections: list[dict],
        viewer: dict | None = None,
    ):
        """Yield the same event sequence as pipelines/find_ask.py, with a
        keyword-driven stand-in for the router. MOCK_DELAY (seconds) paces the
        events."""
        request = (request or "").strip()
        if not request:
            raise ValueError("Tell me what to find.")
        doc_id, info = self._pick_doc(store, doc_ids)
        return self._find_events(
            store, request, doc_id, info, int(top_k), sections or [], viewer or {}
        )

    def _find_events(self, store, request, doc_id, info, top_k, sections, viewer):
        delay = float(os.environ.get("MOCK_DELAY", "0"))
        cur = max(1, int(viewer.get("page") or 1))
        yield {"type": "status", "text": "Looking at your page…"}
        time.sleep(delay)

        kind, a, b = self._mock_route(request, sections)

        if kind == "go_to_section":
            yield {"type": "route", "route": "section", "title": b, "page": a}
            yield {"type": "done", "kind": "navigate", "page": a, "title": b}
            return

        if kind == "point_here":
            target = a
            yield {"type": "route", "route": "local", "target": target}
            yield {"type": "status", "text": "Pinning it down…"}
            time.sleep(delay)
            box = self._mock_box(store, doc_id, cur, target)
            if box is not None:
                yield {"type": "done", "kind": "point", "found": True,
                       "target": target, "page": cur, "bbox": box}
                return
            query = target  # couldn't pin here → fall through to search
        else:
            target, query = a, b

        yield {"type": "route", "route": "search", "target": target, "query": query}
        pages = self._pick_pages(request, info["pages"], top_k)
        images = render_pages(store.pdf_path(doc_id), pages)
        yield {
            "type": "tool_result",
            "tool": "search_docs",
            "gallery": [
                (img, f"{info['name']} — p.{p} (mock)")
                for p, img in zip(pages, images)
            ],
            "page_refs": [(doc_id, p) for p in pages],
        }
        yield {"type": "status", "text": f"Checking {len(pages)} pages…"}
        time.sleep(delay)

        # a magic target exercises the all-NO give-up path
        give_up = any(w in target.lower() for w in ("xyzzy", "flux capacitor", "nonexistent"))
        flags = [False] * len(pages)
        if not give_up and flags:
            flags[0] = True
        yield {
            "type": "classified",
            "results": [[p, f] for p, f in zip(pages, flags)],
        }
        for p, img, f in zip(pages, images, flags):
            if not f:
                continue
            yield {"type": "found", "page": p}
            yield {"type": "status", "text": "Pinning it down…"}
            time.sleep(delay)
            box = self._mock_box(store, doc_id, p, target)
            yield {"type": "done", "kind": "point", "found": True,
                   "target": target, "page": p, "bbox": box}
            return
        yield {"type": "done", "kind": "point", "found": False, "target": target}

    @staticmethod
    def _mock_route(request: str, sections: list[dict]):
        """Fake the LLM router with keyword heuristics. Returns one of:
        ("go_to_section", page, title) / ("point_here", target, None) /
        ("search", target, query)."""
        r = request.lower()
        nav_verb = any(
            v in r for v in ("go to", "take me", "open", "bring up", "pull up",
                             "navigate", "jump to")
        )
        is_circle = any(
            v in r for v in ("circle", "point", "highlight", "mark", "show me where")
        )
        here = any(v in r for v in ("here", "this page", "this", "current"))
        secs = [{"title": o["title"], "page_start": o["page"]} for o in sections]
        best = match_section(request, secs) if secs else None

        if (nav_verb or "section" in r or "chapter" in r) and best and best["score"] >= 0.4:
            return "go_to_section", best["page"], best["title"]
        target = MockAskPipeline._clean_target(request)
        if is_circle and here:
            return "point_here", target, None
        if is_circle:
            return "search", target, request
        if best and best["score"] >= 0.6:
            return "go_to_section", best["page"], best["title"]
        return "search", target, request

    @staticmethod
    def _clean_target(request: str) -> str:
        r = re.sub(
            r"^(?:can you |please )?(?:circle|point (?:at|to)|highlight|mark|"
            r"show me where)\s+",
            "", request.strip(), flags=re.I,
        )
        r = re.sub(r"^(the |a |an )", "", r, flags=re.I)
        r = re.sub(r"\b(?:on |in )?(?:this|the current) page\b", "", r, flags=re.I)
        r = re.sub(r"\bhere\b", "", r, flags=re.I)
        r = re.sub(r"\s+(?:is|are)(?:\s+located)?$", "", r, flags=re.I)
        return " ".join(r.split()) or request

    @staticmethod
    def _mock_box(store: MockStore, doc_id: str, page: int, target: str):
        """A deterministic, target-dependent bbox in rendered-page pixels;
        None for ~1 in 8 targets so the "located but not pinpointed" path is
        exercised too."""
        path = store.pdf_path(doc_id)
        if not path:
            return None
        seed = int(hashlib.sha1(target.encode()).hexdigest(), 16)
        if seed % 8 == 7:
            return None
        w, h = page_size(path, page)
        x1 = w * (0.1 + (seed % 5) * 0.1)
        y1 = h * (0.15 + (seed // 5 % 5) * 0.12)
        return [round(x1), round(y1), round(x1 + w * 0.3), round(y1 + h * 0.18)]

    @staticmethod
    def _pick_doc(store: MockStore, doc_ids: list[str] | None):
        docs = {d["doc_id"]: d for d in store.list_docs()}
        if not docs:
            raise ValueError(
                f"No PDFs in the mock library yet — drop one into {store.pdf_dir}."
            )
        doc_id = (doc_ids or list(docs))[0]
        if doc_id not in docs:
            raise ValueError("That manual isn't in the mock library.")
        return doc_id, docs[doc_id]

    @staticmethod
    def _pick_pages(question: str, n_pages: int, top_k: int) -> list[int]:
        """A deterministic, question-dependent spread of 1-based pages, so
        different questions cite different pages (nice for clicking through the
        viewer) while the same question is stable across reloads."""
        k = max(1, min(top_k, n_pages))
        seed = int(hashlib.sha1(question.encode()).hexdigest(), 16)
        start = seed % n_pages  # 0-based anchor
        step = max(1, n_pages // k)
        out: list[int] = []
        for i in range(k):
            page = (start + i * step) % n_pages + 1  # back to 1-based
            if page not in out:
                out.append(page)
        return out