#!/usr/bin/env python3
"""Chromakey garment-clipping analyzer (T-1089).

Counts key-colour (pure magenta) pixels in each capture PNG produced by
client/tools/garment_qa/chromakey_scene.gd. Any solid patch of key pixels showing
through a garment is a body-clip-through: the QA scene painted the covered body
segments flat magenta, so magenta the camera can see = body poking through cloth.

A capture FAILS when its largest connected key-pixel component is >= --min-pixels
(default 8; a small tolerance for anti-aliased edges). For each failing capture a
highlighted copy is written to <out_dir>/failures/ with the clip pixels recoloured
lime and each component boxed in red. Results land in <out_dir>/report.json plus a
one-screen summary on stdout.

Reads the same JSON config as the capture scene (via --config) to locate out_dir,
or takes --dir directly.
"""

from __future__ import annotations

import argparse
import collections
import json
import sys
from pathlib import Path

from PIL import Image, ImageChops, ImageDraw

# Key-colour gate. Pure magenta is (255, 0, 255); the blue channel is the decisive
# discriminator — skin/garment texture is never simultaneously high-red, low-green
# AND high-blue, so this never fires on legitimate body or cloth pixels.
R_MIN = 200
G_MAX = 60
B_MIN = 200


def build_key_mask(img: Image.Image) -> Image.Image:
    """Return an "L" mask, 255 where the pixel is key-colour, else 0."""
    r, g, b = img.convert("RGB").split()
    r_ok = r.point(lambda v: 255 if v >= R_MIN else 0)
    g_ok = g.point(lambda v: 255 if v <= G_MAX else 0)
    b_ok = b.point(lambda v: 255 if v >= B_MIN else 0)
    return ImageChops.multiply(ImageChops.multiply(r_ok, g_ok), b_ok)


def connected_components(mask: Image.Image) -> list[dict]:
    """8-connected components of the key-pixel mask.

    Only the mask bounding box is scanned, so cost tracks the (sparse) clip area,
    not the whole frame. Returns one dict per component: size + pixel bbox.
    """
    bbox = mask.getbbox()
    if bbox is None:
        return []
    x0, y0, x1, y1 = bbox
    region = mask.crop(bbox)
    width, height = region.size
    px = region.load()

    seen = [[False] * width for _ in range(height)]
    components: list[dict] = []
    for sy in range(height):
        for sx in range(width):
            if px[sx, sy] == 0 or seen[sy][sx]:
                continue
            size = 0
            min_x = max_x = sx
            min_y = max_y = sy
            queue = collections.deque([(sx, sy)])
            seen[sy][sx] = True
            while queue:
                cx, cy = queue.popleft()
                size += 1
                min_x, max_x = min(min_x, cx), max(max_x, cx)
                min_y, max_y = min(min_y, cy), max(max_y, cy)
                for dy in (-1, 0, 1):
                    for dx in (-1, 0, 1):
                        nx, ny = cx + dx, cy + dy
                        if 0 <= nx < width and 0 <= ny < height:
                            if px[nx, ny] != 0 and not seen[ny][nx]:
                                seen[ny][nx] = True
                                queue.append((nx, ny))
            components.append(
                {
                    "size": size,
                    # bbox back in full-image coordinates
                    "bbox": [x0 + min_x, y0 + min_y, x0 + max_x + 1, y0 + max_y + 1],
                }
            )
    components.sort(key=lambda c: c["size"], reverse=True)
    return components


def write_highlight(img: Image.Image, mask: Image.Image, comps: list[dict], dest: Path) -> None:
    """Recolour key pixels lime and box each component in red."""
    out = img.convert("RGB")
    out.paste((0, 255, 0), mask=mask)
    draw = ImageDraw.Draw(out)
    for comp in comps:
        draw.rectangle(comp["bbox"], outline=(255, 0, 0), width=2)
    dest.parent.mkdir(parents=True, exist_ok=True)
    out.save(dest)


def analyze_dir(out_dir: Path, min_pixels: int) -> dict:
    captures = sorted(p for p in out_dir.glob("*.png") if p.is_file())
    fail_dir = out_dir / "failures"
    results: list[dict] = []

    for path in captures:
        img = Image.open(path)
        mask = build_key_mask(img)
        total = mask.histogram()[255]
        comps = connected_components(mask) if total else []
        largest = comps[0]["size"] if comps else 0
        failed = largest >= min_pixels
        if failed:
            write_highlight(img, mask, comps, fail_dir / (path.stem + "_HL.png"))
        results.append(
            {
                "file": path.name,
                "key_pixels": total,
                "components": len(comps),
                "largest_component": largest,
                "component_sizes": [c["size"] for c in comps[:10]],
                "fail": failed,
            }
        )

    return summarize(out_dir, min_pixels, results)


def summarize(out_dir: Path, min_pixels: int, results: list[dict]) -> dict:
    failures = [r for r in results if r["fail"]]
    by_group: dict[str, dict] = {}
    for r in results:
        # filename: <body>__<clip>__f<n>__yaw<deg>.png
        parts = r["file"].split("__")
        group = "__".join(parts[:2]) if len(parts) >= 2 else r["file"]
        entry = by_group.setdefault(group, {"captures": 0, "failures": 0, "worst": 0})
        entry["captures"] += 1
        entry["failures"] += 1 if r["fail"] else 0
        entry["worst"] = max(entry["worst"], r["largest_component"])
    return {
        "out_dir": str(out_dir),
        "min_component_pixels": min_pixels,
        "key_gate": {"r_min": R_MIN, "g_max": G_MAX, "b_min": B_MIN},
        "total_captures": len(results),
        "failures": len(failures),
        "by_group": by_group,
        "captures": results,
    }


def print_summary(report: dict) -> None:
    print("=" * 64)
    print("garment-qa chromakey analysis")
    print(f"  out_dir           : {report['out_dir']}")
    print(f"  min clip pixels   : {report['min_component_pixels']}")
    print(f"  captures          : {report['total_captures']}")
    print(f"  FAILING captures  : {report['failures']}")
    print("-" * 64)
    print(f"  {'group (body__clip)':<28}{'caps':>6}{'fails':>7}{'worst':>7}")
    for group, g in sorted(report["by_group"].items()):
        print(f"  {group:<28}{g['captures']:>6}{g['failures']:>7}{g['worst']:>7}")
    print("=" * 64)


def resolve_out_dir(args: argparse.Namespace) -> Path:
    if args.dir:
        return Path(args.dir)
    if args.config:
        cfg = json.loads(Path(args.config).read_text())
        out = cfg.get("out_dir")
        if not out:
            sys.exit("analyze_captures: config has no 'out_dir'")
        return Path(out)
    sys.exit("analyze_captures: pass --dir or --config")


def main() -> int:
    parser = argparse.ArgumentParser(description="Chromakey garment-clipping analyzer")
    parser.add_argument("--config", help="capture config JSON (reads out_dir from it)")
    parser.add_argument("--dir", help="directory of capture PNGs (overrides --config out_dir)")
    parser.add_argument(
        "--min-pixels",
        type=int,
        default=8,
        help="largest connected key-pixel component that counts as a clip (default 8)",
    )
    parser.add_argument("--report", help="report JSON path (default: <out_dir>/report.json)")
    args = parser.parse_args()

    out_dir = resolve_out_dir(args)
    if not out_dir.is_dir():
        sys.exit(f"analyze_captures: not a directory: {out_dir}")

    report = analyze_dir(out_dir, args.min_pixels)
    report_path = Path(args.report) if args.report else out_dir / "report.json"
    report_path.write_text(json.dumps(report, indent=2))
    print_summary(report)
    print(f"  report : {report_path}")
    if report["failures"]:
        print(f"  highlighted failing frames : {out_dir / 'failures'}")
    return 0


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