#!/usr/bin/env python3
"""
Genera un PDF interactivo (con checkboxes AcroForm) de los diffs reales
de un directorio out_surya*.

Uso:
  python3 build_surya_pdf.py --dir out_surya2 --out out_surya2/report.pdf
"""
import argparse, json, re, sys, io
from pathlib import Path
import fitz  # PyMuPDF

from reportlab.pdfgen import canvas as rlcanvas
from reportlab.lib.pagesizes import A4
from reportlab.lib.units import cm, mm
from reportlab.lib import colors
from reportlab.lib.utils import ImageReader
from reportlab.pdfbase import pdfmetrics
from reportlab.pdfbase.ttfonts import TTFont

HERE = Path(__file__).resolve().parent
CFG  = json.loads((HERE / "books.json").read_text())
SCAN_ROOT = Path(CFG["_scan_root"])

# stem → pdf path
STEM2PDF = {}
for book, rels in CFG["books"].items():
    for rel in rels:
        p = SCAN_ROOT / rel
        if p.exists():
            STEM2PDF.setdefault(book, {})[p.stem] = p

W, H = A4   # 595.28 x 841.89 pt

MARGIN  = 1.8 * cm
COL_L   = MARGIN
COL_R   = W / 2 + 0.4 * cm
COL_W_L = W / 2 - MARGIN - 0.4 * cm
COL_W_R = W / 2 - MARGIN - 0.4 * cm

CONF_COLOR = {
    "high":   colors.HexColor("#2e7d32"),
    "medium": colors.HexColor("#e65100"),
    "low":    colors.HexColor("#b71c1c"),
}

def render_page_jpeg(book, pageref, dpi=110):
    if not pageref or "#" not in pageref:
        return None
    stem, pg = pageref.split("#", 1)
    m = re.search(r'(\d+)', pg)
    if not m:
        return None
    page_num = int(m.group(1))
    pdf_path = STEM2PDF.get(book, {}).get(stem)
    if not pdf_path:
        return None
    try:
        doc  = fitz.open(str(pdf_path))
        if page_num < 1 or page_num > len(doc):
            return None
        page = doc[page_num - 1]
        mat  = fitz.Matrix(dpi / 72, dpi / 72)
        pix  = page.get_pixmap(matrix=mat, colorspace=fitz.csRGB)
        jpeg = pix.tobytes("jpeg", jpg_quality=72)
        doc.close()
        return jpeg
    except Exception as e:
        print(f"  warn: {pdf_path} p{page_num}: {e}", file=sys.stderr)
        return None

def wrap_text(c, text, x, y, max_w, font, size, line_h):
    """Draw wrapped text, return final y."""
    c.setFont(font, size)
    words = text.split()
    line  = ""
    for w in words:
        test = (line + " " + w).strip()
        if c.stringWidth(test, font, size) <= max_w:
            line = test
        else:
            if line:
                c.drawString(x, y, line)
                y -= line_h
            line = w
    if line:
        c.drawString(x, y, line)
        y -= line_h
    return y

def load_rows(out_dir):
    rows = []
    for f in sorted(out_dir.glob("*.adjudicated.jsonl")):
        book = f.stem.replace(".adjudicated", "")
        for line in f.read_text().splitlines():
            e = json.loads(line)
            llm = e.get("llm", {})
            if llm.get("is_real_diff"):
                rows.append({
                    "book":             book,
                    "file":             e.get("file", ""),
                    "page":             e.get("page", ""),
                    "repo":             e.get("repo", ""),
                    "original_reading": llm.get("original_reading", ""),
                    "confidence":       llm.get("confidence", "medium"),
                    "note":             llm.get("note", ""),
                    "context":          e.get("context", ""),
                })
    return rows

def build_pdf(rows, out_path):
    c = rlcanvas.Canvas(str(out_path), pagesize=A4)
    c.setTitle("Surya OCR diffs — correcciones pendientes")

    total = len(rows)
    print(f"Generando PDF con {total} diffs...", file=sys.stderr)

    for idx, r in enumerate(rows):
        # ── Scan image (right column) ──────────────────────────────
        jpeg = render_page_jpeg(r["book"], r["page"])
        if (idx + 1) % 10 == 0:
            print(f"  {idx+1}/{total}", file=sys.stderr, flush=True)

        y_top = H - MARGIN

        # ── Header bar ────────────────────────────────────────────
        c.setFillColor(colors.HexColor("#1a1a2e"))
        c.rect(0, y_top - 1.1*cm, W, 1.1*cm, fill=1, stroke=0)
        c.setFillColor(colors.white)
        c.setFont("Helvetica-Bold", 10)
        short_file = r["file"].split("/")[-1]
        page_label = r["page"].split("#")[-1] if r["page"] else ""
        c.drawString(MARGIN, y_top - 0.75*cm, f"{r['book']}  /  {short_file}  —  {page_label}")

        # diff counter top right
        c.setFont("Helvetica", 8)
        c.drawRightString(W - MARGIN, y_top - 0.75*cm, f"{idx+1} / {total}")

        y = y_top - 1.1*cm - 0.4*cm

        # ── Checkbox (AcroForm) ───────────────────────────────────
        cb_size = 18
        cb_x = COL_L
        cb_y = y - cb_size
        c.acroForm.checkbox(
            name=f"cb_{idx}",
            tooltip=f"Aplicar corrección {idx+1}",
            x=cb_x, y=cb_y,
            size=cb_size,
            buttonStyle="check",
            borderColor=colors.HexColor("#555"),
            fillColor=colors.white,
            textColor=colors.HexColor("#1a1a2e"),
            forceBorder=True,
        )

        # ── Confidence badge ──────────────────────────────────────
        conf      = r["confidence"]
        conf_col  = CONF_COLOR.get(conf, colors.grey)
        badge_x   = cb_x + cb_size + 6
        badge_w   = 52
        c.setFillColor(conf_col)
        c.roundRect(badge_x, cb_y, badge_w, cb_size, 4, fill=1, stroke=0)
        c.setFillColor(colors.white)
        c.setFont("Helvetica-Bold", 8)
        c.drawCentredString(badge_x + badge_w/2, cb_y + 5, conf.upper())

        y = cb_y - 0.5*cm

        # ── Diff line ────────────────────────────────────────────
        c.setFillColor(colors.HexColor("#c62828"))
        c.setFont("Helvetica-Bold", 13)
        repo_w = c.stringWidth(r["repo"], "Helvetica-Bold", 13)
        c.drawString(COL_L, y, r["repo"])

        arrow_x = COL_L + repo_w + 8
        c.setFillColor(colors.HexColor("#888"))
        c.setFont("Helvetica", 13)
        c.drawString(arrow_x, y, "→")
        arrow_w = c.stringWidth("→", "Helvetica", 13)

        c.setFillColor(colors.HexColor("#1b5e20"))
        c.setFont("Helvetica-Bold", 13)
        c.drawString(arrow_x + arrow_w + 8, y, r["original_reading"][:80])

        y -= 0.7*cm

        # ── Context ───────────────────────────────────────────────
        ctx_raw = r["context"]
        ctx_plain = ctx_raw.replace("[[", "").replace("]]", "")
        # highlight span
        m = re.match(r'(.*)\[\[(.+?)\]\](.*)', ctx_raw)
        if m:
            pre, span_text, post = m.group(1), m.group(2), m.group(3)
            c.setFont("Helvetica", 8)
            c.setFillColor(colors.HexColor("#444"))
            pre_w = c.stringWidth(pre, "Helvetica", 8)
            span_w = c.stringWidth(span_text, "Helvetica-Bold", 8)
            # draw pre
            if pre_w < COL_W_L:
                c.drawString(COL_L, y, pre[-80:])
                sx = COL_L + min(pre_w, c.stringWidth(pre[-80:], "Helvetica", 8))
                # highlight box
                c.setFillColor(colors.HexColor("#fff9c4"))
                c.rect(sx - 1, y - 2, span_w + 2, 11, fill=1, stroke=0)
                c.setFillColor(colors.HexColor("#333"))
                c.setFont("Helvetica-Bold", 8)
                c.drawString(sx, y, span_text)
                sx2 = sx + span_w
                c.setFont("Helvetica", 8)
                c.setFillColor(colors.HexColor("#444"))
                c.drawString(sx2, y, post[:60])
            else:
                c.setFillColor(colors.HexColor("#444"))
                c.drawString(COL_L, y, ctx_plain[:120])
        else:
            c.setFont("Helvetica", 8)
            c.setFillColor(colors.HexColor("#444"))
            c.drawString(COL_L, y, ctx_plain[:120])
        y -= 0.55*cm

        # ── Note ─────────────────────────────────────────────────
        c.setFillColor(colors.HexColor("#555"))
        y = wrap_text(c, r["note"], COL_L + 6, y, COL_W_L - 6,
                      "Helvetica-Oblique", 8, 12)

        # ── Separator line ────────────────────────────────────────
        c.setStrokeColor(colors.HexColor("#ddd"))
        c.line(W/2, H - MARGIN - 1.1*cm - 0.2*cm, W/2, MARGIN)

        # ── Scan image (right column) ─────────────────────────────
        if jpeg:
            img    = ImageReader(io.BytesIO(jpeg))
            iw, ih = img.getSize()
            avail_h = H - MARGIN - 1.1*cm - 0.4*cm - MARGIN
            avail_w = COL_W_R
            scale   = min(avail_w / iw, avail_h / ih)
            dw, dh  = iw * scale, ih * scale
            img_x   = COL_R + (avail_w - dw) / 2
            img_y   = H - MARGIN - 1.1*cm - 0.4*cm - dh
            c.drawImage(img, img_x, img_y, dw, dh, preserveAspectRatio=True)
        else:
            c.setFont("Helvetica", 9)
            c.setFillColor(colors.grey)
            c.drawCentredString(COL_R + COL_W_R/2, H/2, "sin imagen")

        # ── Page number footer ────────────────────────────────────
        c.setFont("Helvetica", 7)
        c.setFillColor(colors.HexColor("#aaa"))
        c.drawCentredString(W/2, MARGIN/2, str(idx + 1))

        c.showPage()

    c.save()
    size_mb = Path(out_path).stat().st_size / 1_000_000
    print(f"\nEscrito: {out_path} ({size_mb:.1f} MB, {total} páginas)", file=sys.stderr)

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--dir", default="out_surya2")
    ap.add_argument("--out", default=None)
    a = ap.parse_args()
    out_dir = Path(a.dir)
    out_path = a.out or str(out_dir / "report.pdf")
    rows = load_rows(out_dir)
    build_pdf(rows, out_path)

if __name__ == "__main__":
    main()