"""Step 7: Assemble COCO-format dataset from per-page annotations."""

import json
import logging
import random
import shutil
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Tuple

from PIL import Image

from .config import (
    ANNOTATIONS_DIR,
    COCO_BBOX_PADDING,
    COCO_CATEGORIES,
    IMAGES_DIR,
    PAGE_ANNOTS_DIR,
    TRAIN_FRACTION,
)
from .search import load_documents
from .state import PipelineState

logger = logging.getLogger(__name__)


def _split_documents(
    doc_ids: List[str],
    train_fraction: float = TRAIN_FRACTION,
    seed: int = 42,
) -> Tuple[List[str], List[str]]:
    """Split document IDs into train/val sets."""
    rng = random.Random(seed)
    shuffled = list(doc_ids)
    rng.shuffle(shuffled)
    split_idx = int(len(shuffled) * train_fraction)
    return shuffled[:split_idx], shuffled[split_idx:]


def build_coco(state: PipelineState) -> Dict[str, Any]:
    """Build COCO-format annotation files.

    Returns stats about the built dataset.
    """
    ANNOTATIONS_DIR.mkdir(parents=True, exist_ok=True)

    # Load document metadata for extended fields
    doc_meta = {}
    for doc in load_documents():
        doc_meta[str(doc["id"])] = doc

    # Collect all annotated pages with valid annotations
    doc_pages: Dict[str, List[int]] = {}
    for doc_id in state.all_doc_ids():
        if not state.is_done(doc_id, "validated"):
            continue
        pages = state.get(doc_id, "annotated_pages", [])
        if pages:
            doc_pages[doc_id] = pages

    if not doc_pages:
        logger.warning("No validated documents found. Nothing to build.")
        return {"train_images": 0, "val_images": 0}

    # Split by document
    train_docs, val_docs = _split_documents(list(doc_pages.keys()))
    logger.info("Split: %d train docs, %d val docs", len(train_docs), len(val_docs))

    # Build COCO for each split
    stats = {}
    for split, split_docs in [("train", train_docs), ("val", val_docs)]:
        coco, split_stats = _build_split(split, split_docs, doc_pages, doc_meta)

        # Write COCO JSON
        out_path = ANNOTATIONS_DIR / f"{split}.json"
        with open(out_path, "w") as f:
            json.dump(coco, f, indent=2)

        stats[f"{split}_images"] = split_stats["images"]
        stats[f"{split}_annotations"] = split_stats["annotations"]
        stats[f"{split}_checked"] = split_stats["checked"]
        stats[f"{split}_unchecked"] = split_stats["unchecked"]

        logger.info(
            "%s: %d images, %d annotations (%d checked, %d unchecked)",
            split, split_stats["images"], split_stats["annotations"],
            split_stats["checked"], split_stats["unchecked"],
        )

    return stats


def _build_split(
    split: str,
    doc_ids: List[str],
    doc_pages: Dict[str, List[int]],
    doc_meta: Dict[str, Dict],
) -> Tuple[Dict[str, Any], Dict[str, int]]:
    """Build COCO dict for one split."""
    coco = {
        "info": {
            "description": "Checkbox detection dataset from DocumentCloud government forms",
            "version": "1.0",
            "year": datetime.now().year,
            "date_created": datetime.now().isoformat(),
        },
        "licenses": [],
        "categories": COCO_CATEGORIES,
        "images": [],
        "annotations": [],
    }

    image_id = 0
    annot_id = 0
    stats = {"images": 0, "annotations": 0, "checked": 0, "unchecked": 0}
    split_dir = IMAGES_DIR / split

    for doc_id in doc_ids:
        pages = doc_pages.get(doc_id, [])
        meta = doc_meta.get(doc_id, {})

        for page_idx in pages:
            annot_file = PAGE_ANNOTS_DIR / f"{doc_id}_p{page_idx:03d}.json"
            if not annot_file.exists():
                continue

            with open(annot_file) as f:
                page_annot = json.load(f)

            checkboxes = page_annot.get("checkboxes", [])
            if not checkboxes:
                continue

            # Ensure image is in the correct split directory
            img_filename = f"{doc_id}_p{page_idx:03d}.png"
            target_path = split_dir / img_filename
            if not target_path.exists():
                # Try to find in other split or staging dir
                for other in ("staging", "train", "val"):
                    src = IMAGES_DIR / other / img_filename
                    if src.exists() and src != target_path:
                        split_dir.mkdir(parents=True, exist_ok=True)
                        shutil.copy2(str(src), str(target_path))
                        break

            if not target_path.exists():
                logger.warning("Image not found for %s page %d", doc_id, page_idx)
                continue

            # Get image dimensions
            with Image.open(target_path) as img:
                img_w, img_h = img.size

            image_id += 1
            coco["images"].append({
                "id": image_id,
                "file_name": img_filename,
                "width": img_w,
                "height": img_h,
                "document_id": doc_id,
                "document_title": meta.get("title", ""),
                "source_org": meta.get("organization", ""),
                "page_number": page_idx,
                "source_url": meta.get("canonical_url", ""),
            })
            stats["images"] += 1

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

                # Apply padding and clamp to image bounds
                x0, y0, x1, y1 = bbox
                x0 = max(0, x0 - COCO_BBOX_PADDING)
                y0 = max(0, y0 - COCO_BBOX_PADDING)
                x1 = min(img_w, x1 + COCO_BBOX_PADDING)
                y1 = min(img_h, y1 + COCO_BBOX_PADDING)

                # Convert [x0, y0, x1, y1] to COCO [x, y, w, h]
                coco_bbox = [x0, y0, x1 - x0, y1 - y0]

                is_checked = cb.get("state") == "checked"
                cat_id = 1 if is_checked else 2

                annot_id += 1
                coco["annotations"].append({
                    "id": annot_id,
                    "image_id": image_id,
                    "category_id": cat_id,
                    "bbox": coco_bbox,
                    "area": coco_bbox[2] * coco_bbox[3],
                    "iscrowd": 0,
                    "attributes": {
                        "style": cb.get("style", "square"),
                        "fill_method": cb.get("fill_method", "empty"),
                    },
                })
                stats["annotations"] += 1
                if is_checked:
                    stats["checked"] += 1
                else:
                    stats["unchecked"] += 1

    return coco, stats


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    st = PipelineState()
    build_coco(st)
