"""Step 6: Automated validation / QC for annotations."""

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

from .config import (
    DUPLICATE_IOU_THRESHOLD,
    MAX_BBOX_ASPECT_RATIO,
    MAX_BBOX_SIZE_PX,
    MIN_BBOX_SIZE_PX,
    PAGE_ANNOTS_DIR,
    SPOT_CHECK_FRACTION,
)
from .download import get_pdf_path
from .state import PipelineState

logger = logging.getLogger(__name__)


def _iou(a: List[int], b: List[int]) -> float:
    """Compute IoU between two [x0, y0, x1, y1] bboxes."""
    x0 = max(a[0], b[0])
    y0 = max(a[1], b[1])
    x1 = min(a[2], b[2])
    y1 = min(a[3], b[3])

    inter = max(0, x1 - x0) * max(0, y1 - y0)
    if inter == 0:
        return 0.0

    area_a = (a[2] - a[0]) * (a[3] - a[1])
    area_b = (b[2] - b[0]) * (b[3] - b[1])
    union = area_a + area_b - inter
    return inter / union if union > 0 else 0.0


def validate_annotation(
    annot_data: Dict[str, Any],
    image_width: int,
    image_height: int,
) -> Tuple[List[Dict], List[str]]:
    """Validate a single page annotation.

    Returns (valid_checkboxes, issues).
    """
    issues: List[str] = []
    valid = []
    bboxes_seen = []

    checkboxes = annot_data.get("checkboxes", [])

    for i, cb in enumerate(checkboxes):
        bbox = cb.get("bbox", [])
        if len(bbox) != 4:
            issues.append(f"Checkbox {i}: invalid bbox length {len(bbox)}")
            continue

        x0, y0, x1, y1 = bbox

        # Bounds check
        if x0 < 0 or y0 < 0 or x1 > image_width or y1 > image_height:
            issues.append(f"Checkbox {i}: bbox {bbox} out of image bounds ({image_width}x{image_height})")
            # Clamp rather than discard
            x0 = max(0, min(x0, image_width))
            y0 = max(0, min(y0, image_height))
            x1 = max(0, min(x1, image_width))
            y1 = max(0, min(y1, image_height))
            bbox = [x0, y0, x1, y1]
            cb = {**cb, "bbox": bbox}

        w = x1 - x0
        h = y1 - y0

        # Min size
        if w < MIN_BBOX_SIZE_PX or h < MIN_BBOX_SIZE_PX:
            issues.append(f"Checkbox {i}: too small ({w}x{h}px)")
            continue

        # Max size
        if w > MAX_BBOX_SIZE_PX or h > MAX_BBOX_SIZE_PX:
            issues.append(f"Checkbox {i}: too large ({w}x{h}px)")
            continue

        # Aspect ratio
        aspect = max(w, h) / max(min(w, h), 1)
        if aspect > MAX_BBOX_ASPECT_RATIO:
            issues.append(f"Checkbox {i}: bad aspect ratio {aspect:.1f}")
            continue

        # Duplicate check
        is_dup = False
        for prev_bbox in bboxes_seen:
            if _iou(bbox, prev_bbox) > DUPLICATE_IOU_THRESHOLD:
                issues.append(f"Checkbox {i}: duplicate of existing bbox (IoU > {DUPLICATE_IOU_THRESHOLD})")
                is_dup = True
                break

        if is_dup:
            continue

        bboxes_seen.append(bbox)
        valid.append(cb)

    return valid, issues


def validate(state: PipelineState) -> Dict[str, Any]:
    """Run validation on all annotated pages.

    Returns summary stats.
    """
    stats = {
        "total_pages": 0,
        "pages_with_issues": 0,
        "total_checkboxes_in": 0,
        "total_checkboxes_valid": 0,
        "total_issues": 0,
        "issue_types": {},
    }

    for doc_id in state.all_doc_ids():
        annotated_pages = state.get(doc_id, "annotated_pages", [])
        if not annotated_pages:
            continue

        if state.is_done(doc_id, "validated"):
            continue

        all_valid = True
        for page_idx in annotated_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:
                annot_data = json.load(f)

            # Get image dimensions from the annotation or from image file
            img_width, img_height = _get_image_dimensions(doc_id, page_idx)
            if img_width == 0:
                logger.warning("Cannot determine image size for doc %s page %d", doc_id, page_idx)
                continue

            stats["total_pages"] += 1
            stats["total_checkboxes_in"] += len(annot_data.get("checkboxes", []))

            valid_cbs, issues = validate_annotation(annot_data, img_width, img_height)

            stats["total_checkboxes_valid"] += len(valid_cbs)
            stats["total_issues"] += len(issues)

            if issues:
                stats["pages_with_issues"] += 1
                all_valid = False
                for issue in issues:
                    # Categorize by issue type (first word after "Checkbox N:")
                    parts = issue.split(":")
                    if len(parts) >= 2:
                        issue_type = parts[1].strip().split()[0] if parts[1].strip() else "unknown"
                    else:
                        issue_type = "unknown"
                    stats["issue_types"][issue_type] = stats["issue_types"].get(issue_type, 0) + 1

            # Back up raw Gemini annotation before overwriting
            backup = annot_file.with_suffix(".raw.json")
            if not backup.exists():
                shutil.copy2(annot_file, backup)

            # Save validated annotation back (with only valid checkboxes)
            validated_data = {**annot_data, "checkboxes": valid_cbs}
            with open(annot_file, "w") as f:
                json.dump(validated_data, f, indent=2)

        state.mark_done(doc_id, "validated")

    logger.info(
        "Validation complete. %d pages, %d/%d checkboxes valid, %d issues on %d pages.",
        stats["total_pages"],
        stats["total_checkboxes_valid"],
        stats["total_checkboxes_in"],
        stats["total_issues"],
        stats["pages_with_issues"],
    )
    if stats["issue_types"]:
        logger.info("Issue breakdown: %s", stats["issue_types"])

    return stats


def _get_image_dimensions(doc_id: str, page_idx: int) -> Tuple[int, int]:
    """Get image dimensions for a page, trying train and val directories."""
    from PIL import Image

    from .config import IMAGES_DIR

    for split in ("staging", "train", "val"):
        img_path = IMAGES_DIR / split / f"{doc_id}_p{page_idx:03d}.png"
        if img_path.exists():
            with Image.open(img_path) as img:
                return img.size
    return 0, 0


def cross_engine_spot_check(state: PipelineState, fraction: float = SPOT_CHECK_FRACTION) -> Dict[str, Any]:
    """Run natural-pdf's vector detector on a random sample and compare.

    Returns comparison stats.
    """
    from natural_pdf import PDF

    # Collect all annotated pages
    all_pages = []
    for doc_id in state.all_doc_ids():
        for page_idx in state.get(doc_id, "annotated_pages", []):
            all_pages.append((doc_id, page_idx))

    if not all_pages:
        return {"sample_size": 0}

    sample_size = max(1, int(len(all_pages) * fraction))
    sample = random.sample(all_pages, min(sample_size, len(all_pages)))

    results = {
        "sample_size": len(sample),
        "count_matches": 0,
        "count_mismatches": 0,
        "avg_count_diff": 0.0,
    }

    count_diffs = []

    for doc_id, page_idx in sample:
        # Load Gemini annotation
        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:
            annot_data = json.load(f)
        gemini_count = len(annot_data.get("checkboxes", []))

        # Run vector detector
        pdf_path = get_pdf_path(doc_id)
        if not pdf_path.exists():
            continue

        try:
            pdf = PDF(str(pdf_path))
            try:
                page = pdf.pages[page_idx]
                rects = page.find_all("rect")
                from .structural_filter import is_checkbox_shaped
                vector_count = sum(1 for r in rects if is_checkbox_shaped(r.bbox))
            finally:
                pdf.close()
        except Exception:
            continue

        diff = abs(gemini_count - vector_count)
        count_diffs.append(diff)

        # "Match" if counts are within 50% of each other
        if gemini_count == 0 and vector_count == 0:
            results["count_matches"] += 1
        elif gemini_count > 0 and diff / max(gemini_count, vector_count) < 0.5:
            results["count_matches"] += 1
        else:
            results["count_mismatches"] += 1

    if count_diffs:
        results["avg_count_diff"] = sum(count_diffs) / len(count_diffs)

    logger.info(
        "Spot check: %d/%d count matches, avg diff %.1f",
        results["count_matches"],
        results["sample_size"],
        results["avg_count_diff"],
    )
    return results


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