from __future__ import annotations

import re


_CONTINUED_PREFIX_RE = re.compile(r"^\[continued[^\]]*\]\n")


def chunk_snippet_from_doc_text(doc_text: str, *, max_chars: int = 160) -> str:
    """Recover the original PDF-page text from a stored RAG document's text.

    build_pdf_documents() (src/pdf_manual.py) prepends a synthetic header
    ("OEM Owner's Manual: ...", "Pages: X-Y", "---") that isn't actually in the
    PDF, and prefixes continuation chunks with "[continued p.X-Y]" -- strip
    both so what's left is text that actually appears on the rendered page,
    suitable as a page.search_for() query.
    """
    text = doc_text or ""
    marker = "---\n"
    idx = text.find(marker)
    body = text[idx + len(marker):] if idx >= 0 else text
    body = _CONTINUED_PREFIX_RE.sub("", body)
    return body.strip()[:max_chars]


def _highlight_rects(pg, snippet: str, *, window: int = 5, max_words: int = 90) -> list:
    """Rectangles covering ``snippet`` on page ``pg``.

    The passage text is stored with the PDF's hard line wraps removed, while
    page.search_for() matches within the page's own line structure -- so a long
    verbatim probe often misses. Instead search short word windows (5 words,
    sliding by 4) and union the hits, which survives line wraps and small
    extraction differences and covers the *whole* passage, not just its start.
    """
    words = snippet.split()[:max_words]
    if len(words) < window:
        words_probe = [" ".join(words)] if words else []
    else:
        words_probe = [" ".join(words[i : i + window]) for i in range(0, len(words) - window + 1, window - 1)]
    rects: list = []
    seen: set[tuple] = set()
    for probe in words_probe:
        if len(probe) < 8:
            continue
        for r in pg.search_for(probe):
            key = (round(r.x0), round(r.y0), round(r.x1), round(r.y1))
            if key not in seen:
                seen.add(key)
                rects.append(r)
    return rects


def render_pdf_page_png(
    pdf_path: str,
    *,
    page: int,
    highlight_snippet: str | None = None,
    zoom: float = 2.0,
) -> bytes:
    """Render one PDF page to PNG, best-effort highlighting the given snippet.

    page.search_for() is an exact-substring match against extracted text, so a
    minor whitespace/ligature difference between the stored chunk text and
    fitz's own extraction can make the full snippet miss; shrink it until
    something resolves rather than requiring a verbatim match. No match found
    at all just means a plain, unhighlighted page render -- not an error.
    """
    import fitz  # type: ignore

    doc = fitz.open(pdf_path)
    try:
        page_idx = max(0, min(page - 1, doc.page_count - 1))
        pg = doc.load_page(page_idx)

        if highlight_snippet:
            rects = _highlight_rects(pg, highlight_snippet)
            for r in rects:
                pg.draw_rect(r, color=(0.85, 0.55, 0.1), fill=(1, 0.94, 0.75), fill_opacity=0.4, width=1.4)

        pix = pg.get_pixmap(matrix=fitz.Matrix(zoom, zoom))
        return pix.tobytes("png")
    finally:
        doc.close()