"""Spot-check overlay viewer for annotations."""

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

from PIL import Image, ImageDraw, ImageFont

from .config import IMAGES_DIR, PAGE_ANNOTS_DIR
from .state import PipelineState

logger = logging.getLogger(__name__)


def draw_annotations(
    image: Image.Image,
    checkboxes: List[dict],
    line_width: int = 2,
) -> Image.Image:
    """Overlay checkbox annotations on an image.

    Green = checked, Red = unchecked.
    """
    img = image.copy()
    draw = ImageDraw.Draw(img)

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

        x0, y0, x1, y1 = bbox
        is_checked = cb.get("state") == "checked"
        color = (0, 200, 0) if is_checked else (200, 0, 0)

        draw.rectangle([x0, y0, x1, y1], outline=color, width=line_width)

        # Label
        label = cb.get("state", "?")[0].upper()  # C or U
        style = cb.get("style", "")[0].upper() if cb.get("style") else ""
        text = f"{label}{style}"
        draw.text((x0, max(0, y0 - 12)), text, fill=color)

    return img


def visualize_page(
    doc_id: str,
    page_idx: int,
    split: str = "train",
) -> Optional[Image.Image]:
    """Load image and annotations for one page and return overlay."""
    filename = f"{doc_id}_p{page_idx:03d}.png"
    img_path = IMAGES_DIR / split / filename
    if not img_path.exists():
        # Try other locations
        for other in ("staging", "train", "val"):
            candidate = IMAGES_DIR / other / filename
            if candidate.exists():
                img_path = candidate
                break

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

    annot_path = PAGE_ANNOTS_DIR / f"{doc_id}_p{page_idx:03d}.json"
    if not annot_path.exists():
        logger.warning("Annotation not found for %s page %d", doc_id, page_idx)
        return None

    with open(annot_path) as f:
        annot_data = json.load(f)

    with Image.open(img_path) as im:
        image = im.copy()
    return draw_annotations(image, annot_data.get("checkboxes", []))


def visualize_random(
    state: PipelineState,
    count: int = 10,
    save_dir: Optional[Path] = None,
) -> List[Tuple[str, int, Image.Image]]:
    """Visualize random annotated pages.

    Returns list of (doc_id, page_idx, annotated_image) tuples.
    """
    # 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:
        logger.warning("No annotated pages to visualize")
        return []

    sample = random.sample(all_pages, min(count, len(all_pages)))
    results = []

    for doc_id, page_idx in sample:
        img = visualize_page(doc_id, page_idx)
        if img is not None:
            results.append((doc_id, page_idx, img))

            if save_dir:
                save_dir.mkdir(parents=True, exist_ok=True)
                out_path = save_dir / f"{doc_id}_p{page_idx:03d}_annotated.png"
                img.save(out_path)

    logger.info("Visualized %d/%d pages", len(results), len(sample))
    return results


def visualize_grid(
    state: PipelineState,
    count: int = 16,
    cols: int = 4,
    thumb_width: int = 400,
) -> Optional[Image.Image]:
    """Create a grid of annotated page thumbnails."""
    pages = visualize_random(state, count=count)
    if not pages:
        return None

    thumbs = []
    for _, _, img in pages:
        ratio = thumb_width / img.width
        thumb = img.resize(
            (thumb_width, int(img.height * ratio)),
            Image.Resampling.LANCZOS,
        )
        thumbs.append(thumb)

    rows = (len(thumbs) + cols - 1) // cols
    max_h = max(t.height for t in thumbs)

    grid = Image.new("RGB", (cols * thumb_width, rows * max_h), (255, 255, 255))
    for i, thumb in enumerate(thumbs):
        r, c = divmod(i, cols)
        grid.paste(thumb, (c * thumb_width, r * max_h))

    return grid


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    st = PipelineState()
    results = visualize_random(st, count=5, save_dir=Path("temp/viz"))
    print(f"Saved {len(results)} visualizations to temp/viz/")
