"""
crop_political_articles.py
==========================

Standalone re-cropping pass for political articles in the latest extractor run.

What it does
------------
- Picks the most recent run folder under `output/`.
- For each article folder under `output/<run>/articles/pXXX_aYYY/`
  (these are the political articles — the extractor only writes folders for them),
  reads `info.json` to get the page number and the member region IDs.
- Opens `output/<run>/pages/page_XXX.png` and computes a TIGHT crop:
    union(member region bboxes) with a tiny 5px padding,
    clamped to the page image. This is tighter than the extractor's default
    15px padding around the same union.
- Writes the result to `output/<run>/cropped_articles/pXXX_aYYY.png`.

Guarantees
----------
- Does NOT touch any existing files (articles/, all_political_images/,
  pages/, political_articles.json, political_articles.pdf, etc).
- Only creates a new folder `cropped_articles/` inside the run folder.

Usage
-----
    python crop_political_articles.py
    python crop_political_articles.py --run output/20260528_203315
    python crop_political_articles.py --padding 0
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

from PIL import Image


def find_latest_run(output_root: Path) -> Path:
    if not output_root.exists():
        sys.exit(f"output folder not found: {output_root}")
    runs = sorted(
        [p for p in output_root.iterdir() if p.is_dir() and p.name[0].isdigit()],
        key=lambda p: p.name,
    )
    if not runs:
        sys.exit(f"no run folders found under {output_root}")
    return runs[-1]


def load_regions(pages_dir: Path, page_num: int) -> dict[int, list[int]]:
    """Return {region_id: [x1, y1, x2, y2]} for the given page."""
    regions_path = pages_dir / f"page_{page_num:03d}.regions.json"
    if not regions_path.exists():
        return {}
    with regions_path.open() as f:
        data = json.load(f)
    out: dict[int, list[int]] = {}
    for r in data.get("regions", []):
        rid = r.get("id")
        bbox = r.get("bbox")
        if rid is None or not bbox or len(bbox) != 4:
            continue
        out[int(rid)] = [int(v) for v in bbox]
    return out


def tight_bbox(member_ids: list[int], region_map: dict[int, list[int]],
               fallback_bbox: list[int] | None) -> list[int] | None:
    """Compute union bbox of member regions; fall back to info.json bbox."""
    boxes = [region_map[rid] for rid in member_ids if rid in region_map]
    if boxes:
        x1 = min(b[0] for b in boxes)
        y1 = min(b[1] for b in boxes)
        x2 = max(b[2] for b in boxes)
        y2 = max(b[3] for b in boxes)
        return [x1, y1, x2, y2]
    if fallback_bbox and len(fallback_bbox) == 4:
        return [int(v) for v in fallback_bbox]
    return None


def crop_one(article_dir: Path, pages_dir: Path, out_dir: Path,
             padding: int) -> tuple[bool, str]:
    info_path = article_dir / "info.json"
    if not info_path.exists():
        return False, f"{article_dir.name}: no info.json"
    with info_path.open() as f:
        info = json.load(f)

    page_num = int(info.get("page", 0))
    member_ids = [int(x) for x in info.get("member_region_ids", []) or []]
    fallback = info.get("bbox")

    page_img_path = pages_dir / f"page_{page_num:03d}.png"
    if not page_img_path.exists():
        return False, f"{article_dir.name}: page image missing ({page_img_path.name})"

    region_map = load_regions(pages_dir, page_num)
    bbox = tight_bbox(member_ids, region_map, fallback)
    if bbox is None:
        return False, f"{article_dir.name}: could not compute bbox"

    with Image.open(page_img_path) as im:
        pw, ph = im.size
        x1 = max(0, bbox[0] - padding)
        y1 = max(0, bbox[1] - padding)
        x2 = min(pw, bbox[2] + padding)
        y2 = min(ph, bbox[3] + padding)
        if x2 <= x1 or y2 <= y1:
            return False, f"{article_dir.name}: empty crop {[x1,y1,x2,y2]}"
        cropped = im.crop((x1, y1, x2, y2))
        out_path = out_dir / f"{article_dir.name}.png"
        cropped.save(out_path)
    return True, f"{article_dir.name}: {x2-x1}x{y2-y1}px -> {out_path.name}"


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument(
        "--run", type=Path, default=None,
        help="Path to a specific run folder. Default: latest under ./output",
    )
    ap.add_argument(
        "--output-root", type=Path, default=Path(__file__).parent / "output",
        help="Root output dir (default: ./output next to this script)",
    )
    ap.add_argument(
        "--padding", type=int, default=5,
        help="Pixels of padding around the tight region union (default: 5)",
    )
    args = ap.parse_args()

    run_dir = args.run if args.run is not None else find_latest_run(args.output_root)
    run_dir = run_dir.resolve()
    print(f"Run folder: {run_dir}")

    articles_dir = run_dir / "articles"
    pages_dir = run_dir / "pages"
    if not articles_dir.is_dir():
        sys.exit(f"missing articles/ in {run_dir}")
    if not pages_dir.is_dir():
        sys.exit(f"missing pages/ in {run_dir}")

    out_dir = run_dir / "cropped_articles"
    out_dir.mkdir(exist_ok=True)
    print(f"Writing to: {out_dir}")

    article_folders = sorted(p for p in articles_dir.iterdir() if p.is_dir())
    if not article_folders:
        print("No political article folders found.")
        return 0

    ok = 0
    for adir in article_folders:
        success, msg = crop_one(adir, pages_dir, out_dir, args.padding)
        print(("  ok " if success else "  -- ") + msg)
        if success:
            ok += 1

    print(f"\nDone. {ok}/{len(article_folders)} articles cropped into "
          f"{out_dir.relative_to(run_dir.parent)}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())