"""Step 1: Extract check mark template crops from annotated pages.

Scans page_annotations/ for checked checkboxes and crops the corresponding
region (with 10px padding) from the rendered staging image.  Saves crops
to templates/raw/ and writes a manifest.json listing each crop with metadata.

After running, the user should manually curate:
  - Delete bad crops (e.g., non-checkbox regions, too blurry)
  - Trim to actual check mark (removing excess whitespace)
  - Copy approved templates to templates/curated/

Usage:
    python -m datasets.checkbox.tools.extract_check_templates
    python datasets/checkbox/tools/extract_check_templates.py
"""

import json
import logging
import sys
from pathlib import Path

from PIL import Image

logger = logging.getLogger(__name__)

ROOT = Path(__file__).resolve().parent.parent  # datasets/checkbox/
PAGE_ANNOTS_DIR = ROOT / "raw" / "page_annotations"
STAGING_DIR = ROOT / "images" / "staging"
PDFS_DIR = ROOT / "raw" / "pdfs"
RAW_TEMPLATES_DIR = ROOT / "templates" / "raw"

PADDING = 10  # Extra pixels on each side for user re-cropping
MAX_PER_DOC = 2  # Cap templates per source document for diversity
RENDER_DPI = 150  # Match annotate.py rendering
MAX_IMAGE_LONGEST_SIDE = 1000  # Match annotate.py resize


def _render_staging_image(doc_id: str, page_idx: int) -> Image.Image | None:
    """Render a staging image from PDF when the pre-rendered one is missing."""
    pdf_path = PDFS_DIR / f"{doc_id}.pdf"
    if not pdf_path.exists():
        return None

    try:
        from natural_pdf import PDF
    except ImportError:
        logger.warning("natural_pdf not available for rendering")
        return None

    try:
        pdf = PDF(str(pdf_path))
        if page_idx >= len(pdf.pages):
            pdf.close()
            return None

        page = pdf.pages[page_idx]
        try:
            pil_image = page.to_image(resolution=RENDER_DPI)
        except Exception:
            pil_image = page.render(resolution=RENDER_DPI)

        # Resize so longest side = MAX_IMAGE_LONGEST_SIDE (matches annotate.py)
        orig_w, orig_h = pil_image.size
        longest = max(orig_w, orig_h)
        if longest > MAX_IMAGE_LONGEST_SIDE:
            scale = MAX_IMAGE_LONGEST_SIDE / longest
            pil_image = pil_image.resize((int(orig_w * scale), int(orig_h * scale)))

        # Save to staging for future use
        img_path = STAGING_DIR / f"{doc_id}_p{page_idx:03d}.png"
        STAGING_DIR.mkdir(parents=True, exist_ok=True)
        pil_image.save(img_path, format="PNG")
        logger.info("Rendered missing staging image: %s", img_path.name)

        pdf.close()
        return pil_image
    except Exception as e:
        logger.warning("Failed to render %s page %d: %s", doc_id, page_idx, e)
        return None


def extract_templates():
    """Extract checked checkbox crops from staging images."""
    RAW_TEMPLATES_DIR.mkdir(parents=True, exist_ok=True)

    # Clean old crops
    for old in RAW_TEMPLATES_DIR.glob("check_*.png"):
        old.unlink()

    # Find all annotation files (skip .raw.json backups)
    ann_files = sorted(PAGE_ANNOTS_DIR.glob("*.json"))
    ann_files = [f for f in ann_files if not f.name.endswith(".raw.json")]

    manifest = []
    crop_idx = 0
    rendered = 0
    doc_counts = {}  # doc_id → number of crops taken

    for ann_file in ann_files:
        # Parse doc_id and page_idx from filename: {doc_id}_p{idx:03d}.json
        parts = ann_file.stem.rsplit("_p", 1)
        if len(parts) != 2:
            continue
        doc_id, page_str = parts
        page_idx = int(page_str)

        # Skip if we already have enough from this document
        if doc_counts.get(doc_id, 0) >= MAX_PER_DOC:
            continue

        with open(ann_file) as f:
            data = json.load(f)

        checkboxes = data.get("checkboxes", [])
        checked = [cb for cb in checkboxes if cb.get("state") == "checked"]
        if not checked:
            continue

        # Load staging image, rendering from PDF if missing
        img_path = STAGING_DIR / f"{doc_id}_p{page_idx:03d}.png"
        if img_path.exists():
            img = Image.open(img_path)
        else:
            img = _render_staging_image(doc_id, page_idx)
            if img is None:
                logger.warning("Cannot get image for %s page %d", doc_id, page_idx)
                continue
            rendered += 1

        img_w, img_h = img.size

        for cb in checked:
            if doc_counts.get(doc_id, 0) >= MAX_PER_DOC:
                break

            bbox = cb.get("bbox", [])
            if len(bbox) != 4:
                continue

            x0, y0, x1, y1 = bbox

            # Add padding, clamped to image bounds
            px0 = max(0, x0 - PADDING)
            py0 = max(0, y0 - PADDING)
            px1 = min(img_w, x1 + PADDING)
            py1 = min(img_h, y1 + PADDING)

            # Skip tiny crops
            if (px1 - px0) < 5 or (py1 - py0) < 5:
                continue

            crop = img.crop((px0, py0, px1, py1))

            crop_name = f"check_{crop_idx:04d}.png"
            crop.save(RAW_TEMPLATES_DIR / crop_name)

            manifest.append({
                "filename": crop_name,
                "source_doc": doc_id,
                "page_idx": page_idx,
                "original_bbox": bbox,
                "padded_bbox": [px0, py0, px1, py1],
                "crop_size": [px1 - px0, py1 - py0],
                "staging_image": img_path.name,
                "style": cb.get("style", ""),
                "fill_method": cb.get("fill_method", ""),
            })
            crop_idx += 1
            doc_counts[doc_id] = doc_counts.get(doc_id, 0) + 1

        img.close()

    # Write manifest
    manifest_path = RAW_TEMPLATES_DIR / "manifest.json"
    with open(manifest_path, "w") as f:
        json.dump(manifest, f, indent=2)

    logger.info(
        "Extracted %d checked checkbox crops from %d annotation files "
        "(%d staging images rendered from PDF)",
        crop_idx, len(ann_files), rendered,
    )
    logger.info("Raw templates saved to: %s", RAW_TEMPLATES_DIR)
    logger.info("Manifest: %s", manifest_path)
    logger.info(
        "\nNext steps:\n"
        "  1. Review crops in templates/raw/\n"
        "  2. Delete bad ones, trim to actual check marks\n"
        "  3. Copy approved templates to templates/curated/"
    )

    return manifest


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
    extract_templates()
