"""Extract a PCI template from a frame.

Uses HSV green masking to find the PCI cluster near frame center, then crops

a bounding box around it and saves it as configs/templates/pci_circle.png

so pci_tracker.py will auto-switch to template matching.

Also saves the HSV mask as logs/pci_mask_<frame>.png so you can see what the

green threshold is picking up.

"""

from __future__ import annotations

import argparse

import sys

from pathlib import Path

import cv2

import numpy as np

from rich.console import Console

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from cv._common import load_config

console = Console()

ROOT = Path(__file__).resolve().parents[1]

TEMPLATE_DIR = ROOT / "configs" / "templates"

LOG_DIR = ROOT / "logs"

def main() -> int:

    ap = argparse.ArgumentParser(description="Extract PCI template from a frame.")

    ap.add_argument("frame_path", help="Path to a captured frame PNG.")

    ap.add_argument("--pad", type=int, default=10, help="Pixel margin around bbox.")

    ap.add_argument("--center-crop-frac", type=float, default=0.6,

                    help="Restrict search to the central fraction of the frame (0..1) to avoid UI green.")

    args = ap.parse_args()

    frame_path = Path(args.frame_path)

    frame = cv2.imread(str(frame_path), cv2.IMREAD_COLOR)

    if frame is None:

        console.print(f"[red]Could not read {frame_path}[/red]")

        return 2

    cfg = load_config()

    pci_cfg = cfg["cv"]["pci"]

    hsv_low = np.array(pci_cfg["hsv_low"], dtype=np.uint8)

    hsv_high = np.array(pci_cfg["hsv_high"], dtype=np.uint8)

    h, w = frame.shape[:2]

    hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)

    mask = cv2.inRange(hsv, hsv_low, hsv_high)

    cf = float(args.center_crop_frac)

    cw = int(w * cf); ch = int(h * cf)

    x0 = (w - cw) // 2; y0 = (h - ch) // 2

    center_mask = np.zeros_like(mask)

    center_mask[y0:y0 + ch, x0:x0 + cw] = 255

    mask_c = cv2.bitwise_and(mask, center_mask)

    mask_c = cv2.morphologyEx(mask_c, cv2.MORPH_CLOSE, np.ones((15, 15), np.uint8), iterations=2)

    mask_c = cv2.morphologyEx(mask_c, cv2.MORPH_OPEN, np.ones((3, 3), np.uint8), iterations=1)

    ys, xs = np.where(mask_c > 0)

    if len(xs) == 0:

        console.print("[red]No green pixels found in central region - check HSV range or center-crop.[/red]")

        mask_out = LOG_DIR / f"pci_mask_{frame_path.stem}.png"

        cv2.imwrite(str(mask_out), mask_c)

        console.print(f"[yellow]Saved (empty) mask -> {mask_out}[/yellow]")

        return 2

    x_min, x_max = int(xs.min()), int(xs.max())

    y_min, y_max = int(ys.min()), int(ys.max())

    x_min = max(0, x_min - args.pad)

    y_min = max(0, y_min - args.pad)

    x_max = min(w - 1, x_max + args.pad)

    y_max = min(h - 1, y_max + args.pad)

    console.print(f"PCI bbox: x=({x_min}..{x_max})  y=({y_min}..{y_max})  "

                  f"w={x_max-x_min}  h={y_max-y_min}  green_px={len(xs)}")

    crop = frame[y_min:y_max + 1, x_min:x_max + 1].copy()

    TEMPLATE_DIR.mkdir(parents=True, exist_ok=True)

    tpl_out = TEMPLATE_DIR / "pci_circle.png"

    cv2.imwrite(str(tpl_out), crop)

    console.print(f"[green]Saved template -> {tpl_out}  ({crop.shape[1]}x{crop.shape[0]})[/green]")

    ann = frame.copy()

    cv2.rectangle(ann, (x_min, y_min), (x_max, y_max), (0, 255, 0), 2)

    ann_out = LOG_DIR / f"pci_extract_{frame_path.stem}.png"

    cv2.imwrite(str(ann_out), ann)

    mask_out = LOG_DIR / f"pci_mask_{frame_path.stem}.png"

    cv2.imwrite(str(mask_out), mask_c)

    console.print(f"[green]Saved annotated -> {ann_out}[/green]")

    console.print(f"[green]Saved mask      -> {mask_out}[/green]")

    return 0

if __name__ == "__main__":

    sys.exit(main())