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

The capture scene writes two PNGs per view at the identical (paused) animation frame:
  <stem>.png        pass A — covered body segments flat magenta, garment normal.
  <stem>__shift.png pass B — same, but the garment is flat CYAN and nudged
                    CLIP_EPSILON metres toward the camera (a view-space depth bias).

A body-key pixel that is magenta in A but CYAN in B means the epsilon-shifted garment
now covers it — the body sat within epsilon in front of the cloth. Intersecting the two
separates depth-proximate clip-through from mere overlap:
  clip_through — magenta in A AND cyan in B: skin <= epsilon in front of cloth = poking
                 through. The true defect. Largest connected blob >= --min-pixels
                 (default 8, tolerating AA edges) FAILS the capture.
  exposed_skin — magenta in A AND still magenta in B: skin well in front of any cloth
                 (a limb crossing the torso, an open collar, a bare arm over background).
                 Informational only; does not gate.

For each capture a highlighted copy lands in <out_dir>/failures/ with exposed_skin
recoloured lime and clip_through filled red (+ red box per clip blob). 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 (pass A). 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.
KEY_R_MIN = 200
KEY_G_MAX = 60
KEY_B_MIN = 200

# Shifted-garment gate (pass B). Flat cyan is (0, 255, 255); low red + high green/blue
# isolates the shifted cloth from magenta body-key, skin, and the dark background.
CYAN_R_MAX = 60
CYAN_G_MIN = 190
CYAN_B_MIN = 190


def _threshold(band: Image.Image, lo: int | None, hi: int | None) -> Image.Image:
    def f(v: int) -> int:
        if lo is not None and v < lo:
            return 0
        if hi is not None and v > hi:
            return 0
        return 255

    return band.point(f)


def build_key_mask(img: Image.Image) -> Image.Image:
    """"L" mask, 255 where the pass-A pixel is key-colour (magenta), else 0."""
    r, g, b = img.convert("RGB").split()
    r_ok = _threshold(r, KEY_R_MIN, None)
    g_ok = _threshold(g, None, KEY_G_MAX)
    b_ok = _threshold(b, KEY_B_MIN, None)
    return ImageChops.multiply(ImageChops.multiply(r_ok, g_ok), b_ok)


def build_cyan_mask(img: Image.Image) -> Image.Image:
    """"L" mask, 255 where the pass-B pixel is the shifted flat-cyan garment, else 0."""
    r, g, b = img.convert("RGB").split()
    r_ok = _threshold(r, None, CYAN_R_MAX)
    g_ok = _threshold(g, CYAN_G_MIN, None)
    b_ok = _threshold(b, CYAN_B_MIN, None)
    return ImageChops.multiply(ImageChops.multiply(r_ok, g_ok), b_ok)


def connected_components(mask: Image.Image) -> list[dict]:
    """8-connected components of a 0/255 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": [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, exposed: Image.Image, clip: Image.Image, comps: list[dict], dest: Path
) -> None:
    """Recolour exposed_skin lime, fill clip_through red, box each clip blob."""
    out = img.convert("RGB")
    out.paste((0, 255, 0), mask=exposed)
    out.paste((255, 0, 0), mask=clip)
    draw = ImageDraw.Draw(out)
    for comp in comps:
        draw.rectangle(comp["bbox"], outline=(255, 80, 80), 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() and not p.stem.endswith("__shift")
    )
    fail_dir = out_dir / "failures"
    results: list[dict] = []

    for path in captures:
        shift_path = path.with_name(path.stem + "__shift.png")
        img = Image.open(path)
        key = build_key_mask(img)

        if shift_path.exists():
            cyan = build_cyan_mask(Image.open(shift_path))
            clip = ImageChops.multiply(key, cyan)
            exposed = ImageChops.multiply(key, ImageChops.invert(cyan))
            shift_missing = False
        else:
            # No shift pass — cannot classify; treat all key as exposed, clip empty.
            clip = Image.new("L", img.size, 0)
            exposed = key
            shift_missing = True

        comps = connected_components(clip)
        largest = comps[0]["size"] if comps else 0
        failed = largest >= min_pixels
        if failed or exposed.getbbox() is not None:
            write_highlight(img, exposed, clip, comps, fail_dir / (path.stem + "_HL.png"))
        results.append(
            {
                "file": path.name,
                "key_pixels": key.histogram()[255],
                "exposed_skin_pixels": exposed.histogram()[255],
                "clip_through_pixels": clip.histogram()[255],
                "clip_components": len(comps),
                "largest_clip_component": largest,
                "clip_component_sizes": [c["size"] for c in comps[:10]],
                "fail": failed,
                "shift_pass_missing": shift_missing,
            }
        )

    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, "clip_failures": 0, "worst_clip": 0, "worst_exposed": 0}
        )
        entry["captures"] += 1
        entry["clip_failures"] += 1 if r["fail"] else 0
        entry["worst_clip"] = max(entry["worst_clip"], r["largest_clip_component"])
        entry["worst_exposed"] = max(entry["worst_exposed"], r["exposed_skin_pixels"])
    return {
        "out_dir": str(out_dir),
        "min_component_pixels": min_pixels,
        "key_gate": {"r_min": KEY_R_MIN, "g_max": KEY_G_MAX, "b_min": KEY_B_MIN},
        "shift_gate": {"r_max": CYAN_R_MAX, "g_min": CYAN_G_MIN, "b_min": CYAN_B_MIN},
        "total_captures": len(results),
        "clip_through_failures": len(failures),
        "by_group": by_group,
        "captures": results,
    }


def print_summary(report: dict) -> None:
    print("=" * 74)
    print("garment-qa chromakey analysis (two-pass: clip_through gates, exposed_skin info)")
    print(f"  out_dir              : {report['out_dir']}")
    print(f"  min clip blob pixels : {report['min_component_pixels']}")
    print(f"  captures             : {report['total_captures']}")
    print(f"  CLIP-THROUGH failures: {report['clip_through_failures']}")
    print("-" * 74)
    hdr = f"  {'group (body__clip)':<26}{'caps':>6}{'clipfail':>10}{'worstClip':>11}{'worstExp':>10}"
    print(hdr)
    for group, g in sorted(report["by_group"].items()):
        print(
            f"  {group:<26}{g['captures']:>6}{g['clip_failures']:>10}"
            f"{g['worst_clip']:>11}{g['worst_exposed']:>10}"
        )
    print("=" * 74)


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 clip-through blob that counts as a failure (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["clip_through_failures"]:
        print(f"  highlighted frames : {out_dir / 'failures'}")
    return 0


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