"""Step 2: Download a subset of CommonForms for checkbox location data.

Downloads images and checkbox location annotations from the CommonForms
dataset on HuggingFace (jbarrow/CommonForms).  Each page has bounding box
annotations for checkboxes (single class: 'choice_button', no checked/unchecked
distinction).

Saves:
  - commonforms/images/{split}_{idx:06d}.png  (page images)
  - commonforms/annotations/{split}_{idx:06d}.json  (bbox annotations)
  - commonforms/metadata.json  (dataset stats)

Usage:
    python -m datasets.checkbox.tools.download_commonforms
    python datasets/checkbox/tools/download_commonforms.py [--max-pages 5000]
"""

import argparse
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/
COMMONFORMS_DIR = ROOT / "commonforms"
IMAGES_DIR = COMMONFORMS_DIR / "images"
ANNOTATIONS_DIR = COMMONFORMS_DIR / "annotations"

# CommonForms uses COCO categories; choice_button is the checkbox class.
CHOICE_BUTTON_CATEGORY = "choice_button"


def download_commonforms(max_pages: int = 5000, min_checkboxes: int = 2):
    """Download CommonForms pages that contain checkboxes.

    Args:
        max_pages: Maximum number of pages to download.
        min_checkboxes: Minimum number of choice_button annotations per page.
    """
    try:
        from datasets import load_dataset
    except ImportError:
        raise ImportError(
            "HuggingFace datasets is required. "
            "Install with: pip install datasets"
        )

    IMAGES_DIR.mkdir(parents=True, exist_ok=True)
    ANNOTATIONS_DIR.mkdir(parents=True, exist_ok=True)

    logger.info("Loading CommonForms dataset from HuggingFace...")
    # Load with streaming to avoid downloading the full dataset at once
    ds = load_dataset("jbarrow/CommonForms", split="train", streaming=True)

    # Build category lookup from first example
    category_map = None  # Will be set from first example
    target_cat_ids = set()

    saved = 0
    skipped = 0
    total_checkboxes = 0

    for idx, example in enumerate(ds):
        if saved >= max_pages:
            break

        # Extract COCO-style annotations from the example
        # CommonForms format: each example has 'image', 'objects' with
        # 'bbox', 'category' fields
        objects = example.get("objects", {})
        image = example.get("image")

        if image is None or not objects:
            skipped += 1
            continue

        # CommonForms objects format:
        # {'bbox': [[x0, y0, w, h], ...], 'category': [cat_id, ...]}
        bboxes = objects.get("bbox", [])
        categories = objects.get("category", [])

        if not bboxes or not categories:
            skipped += 1
            continue

        # On first example, build category map if available
        if category_map is None:
            # CommonForms category IDs:
            # 0 = textbox, 1 = choice_button, 2 = signature
            # We target category 1 (choice_button)
            target_cat_ids = {1}
            category_map = {0: "textbox", 1: "choice_button", 2: "signature"}
            logger.info("Category map: %s", category_map)
            logger.info("Targeting categories: %s", target_cat_ids)

        # Filter to checkbox annotations only
        checkbox_bboxes = []
        for bbox, cat in zip(bboxes, categories):
            if cat in target_cat_ids:
                checkbox_bboxes.append(bbox)

        if len(checkbox_bboxes) < min_checkboxes:
            skipped += 1
            continue

        # Save image
        img_name = f"cf_{saved:06d}.png"
        img_path = IMAGES_DIR / img_name

        if isinstance(image, Image.Image):
            img = image.convert("RGB")
        else:
            # If it's a path or bytes, convert
            img = Image.open(image).convert("RGB")

        img.save(img_path)
        img_w, img_h = img.size

        # Save annotation (convert to our format: [x0, y0, x1, y1])
        ann_data = {
            "source": "CommonForms",
            "source_index": idx,
            "image_file": img_name,
            "image_size": [img_w, img_h],
            "checkboxes": [],
        }

        for bbox in checkbox_bboxes:
            # CommonForms bbox format: [x0, y0, width, height] (COCO-style)
            x0, y0, w, h = bbox
            ann_data["checkboxes"].append({
                "bbox": [int(x0), int(y0), int(x0 + w), int(y0 + h)],
                "state": "unknown",  # CommonForms has no checked/unchecked labels
            })

        ann_name = f"cf_{saved:06d}.json"
        with open(ANNOTATIONS_DIR / ann_name, "w") as f:
            json.dump(ann_data, f, indent=2)

        total_checkboxes += len(checkbox_bboxes)
        saved += 1

        if saved % 500 == 0:
            logger.info(
                "Progress: %d pages saved (%d skipped), %d total checkboxes",
                saved, skipped, total_checkboxes,
            )

    # Write metadata
    metadata = {
        "total_pages": saved,
        "total_skipped": skipped,
        "total_checkboxes": total_checkboxes,
        "avg_checkboxes_per_page": total_checkboxes / saved if saved else 0,
        "min_checkboxes_filter": min_checkboxes,
        "max_pages_requested": max_pages,
        "category_map": category_map,
    }

    with open(COMMONFORMS_DIR / "metadata.json", "w") as f:
        json.dump(metadata, f, indent=2)

    logger.info(
        "\nDone: %d pages saved, %d checkboxes total (avg %.1f/page)",
        saved, total_checkboxes, metadata["avg_checkboxes_per_page"],
    )
    logger.info("Images: %s", IMAGES_DIR)
    logger.info("Annotations: %s", ANNOTATIONS_DIR)

    return metadata


def main():
    parser = argparse.ArgumentParser(
        description="Download CommonForms checkbox subset"
    )
    parser.add_argument(
        "--max-pages", type=int, default=5000,
        help="Maximum pages to download (default: 5000)",
    )
    parser.add_argument(
        "--min-checkboxes", type=int, default=2,
        help="Min checkboxes per page to include (default: 2)",
    )
    args = parser.parse_args()

    download_commonforms(
        max_pages=args.max_pages,
        min_checkboxes=args.min_checkboxes,
    )


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