"""
Quick YOLO person segmentation smoke test.

Usage:
  python -m tools.test_yolo_persons <image_path>

Reads YOLO_* env vars:
  YOLO_MODEL_PATH (optional absolute/relative path to .pt)
  YOLO_MODEL (default: yolov8n-seg.pt)
  YOLO_DEVICE (e.g., "cpu" or "0")
  YOLO_CONF (default: 0.25)

Outputs:
  - Prints detection summary to stdout
  - Saves masks to artifacts/tmp/yolo_debug/<image_stem>/ as PNGs
"""
from __future__ import annotations

import os
import sys
from pathlib import Path
import argparse


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("image", help="Path to input image")
    args = ap.parse_args()

    ipath = Path(args.image)
    if not ipath.exists():
        print(f"error: image not found: {ipath}")
        return 2

    try:
        from PIL import Image, ImageDraw  # type: ignore
    except Exception as e:
        print(f"error: Pillow not available: {e}")
        return 3
    try:
        import torch  # type: ignore
        from ultralytics import YOLO  # type: ignore
        # Ensure YOLO weights are downloaded to repo models/
        try:
            from ultralytics import settings as _ul_settings  # type: ignore
            repo_root = Path(__file__).resolve().parents[1]
            weights_dir = repo_root / "models"
            weights_dir.mkdir(parents=True, exist_ok=True)
            _ul_settings.update({"weights_dir": str(weights_dir)})
        except Exception:
            pass
    except Exception as e:
        print(f"error: ultralytics import failed: {e}")
        return 4

    model_name = os.getenv("YOLO_MODEL", "yolov8n-seg.pt")
    model_path = os.getenv("YOLO_MODEL_PATH")
    # If only a model name is given, resolve to models/<name>.pt so downloads go there
    if not model_path:
        try:
            repo_root = Path(__file__).resolve().parents[1]
            weights_dir = repo_root / "models"
            weights_dir.mkdir(parents=True, exist_ok=True)
            model_path = str((weights_dir / Path(str(model_name)).name).resolve())
        except Exception:
            pass
    device = os.getenv("YOLO_DEVICE")
    try:
        conf = float(os.getenv("YOLO_CONF", "0.25"))
    except Exception:
        conf = 0.25

    print(f"torch: {torch.__version__} cuda_available={torch.cuda.is_available()}")
    print(f"model_name={model_name} model_path={model_path} device={device} conf={conf}")

    try:
        model = YOLO(model_path or model_name)
        print("loaded segmentation model")
    except Exception as e:
        print(f"segmentation model load failed: {e}")
        try:
            try:
                repo_root = Path(__file__).resolve().parents[1]
                weights_dir = repo_root / "models"
                weights_dir.mkdir(parents=True, exist_ok=True)
                det_path = str((weights_dir / "yolov8n.pt").resolve())
            except Exception:
                det_path = "yolov8n.pt"
            model = YOLO(det_path)
            print("loaded detection model yolov8n.pt (bbox fallback)")
        except Exception as e2:
            print(f"detection model load failed: {e2}")
            return 5

    # Read size for imgsz heuristic
    try:
        with Image.open(ipath) as im:
            w, h = im.size
    except Exception as e:
        print(f"error: failed to open image: {e}")
        return 6

    try:
        results = model.predict(source=str(ipath), classes=[0], imgsz=max(w, h), conf=conf, verbose=False, device=device)
    except Exception as e:
        print(f"predict error at conf={conf}: {e}")
        try:
            low_conf = max(0.05, conf * 0.5)
            print(f"retrying with conf={low_conf}")
            results = model.predict(source=str(ipath), classes=[0], imgsz=max(w, h), conf=low_conf, verbose=False, device=device)
        except Exception as e2:
            print(f"predict retry failed: {e2}")
            return 7

    save_root = Path("artifacts") / "tmp" / "yolo_debug" / ipath.stem
    save_root.mkdir(parents=True, exist_ok=True)

    persons = 0
    for i, r in enumerate(results):
        try:
            if getattr(r, "masks", None) is not None and r.masks is not None and len(r.masks) > 0:
                print(f"batch {i}: masks={len(r.masks)} boxes={len(r.boxes)}")
                for j, (seg, box, score, cls) in enumerate(zip(r.masks.data, r.boxes.xyxy, r.boxes.conf, r.boxes.cls), start=1):
                    m = (seg.cpu().numpy() > 0.5).astype("uint8") * 255
                    pil = Image.fromarray(m)
                    out = save_root / f"{ipath.stem}_person{j:02d}_mask.png"
                    pil.save(out)
                    x1, y1, x2, y2 = [int(v) for v in box.tolist()]
                    print(f"  person{j:02d}: score={float(score):.3f} cls={int(cls)} box=({x1},{y1},{x2},{y2}) -> {out}")
                    persons += 1
            else:
                print(f"batch {i}: no seg masks; using bbox fallback with {len(r.boxes)} boxes")
                for j, (box, score, cls) in enumerate(zip(r.boxes.xyxy, r.boxes.conf, r.boxes.cls), start=1):
                    x1, y1, x2, y2 = [int(v) for v in box.tolist()]
                    pil = Image.new("L", (w, h), 0)
                    d = ImageDraw.Draw(pil)
                    d.rectangle([(x1, y1), (x2, y2)], fill=255)
                    out = save_root / f"{ipath.stem}_person{j:02d}_bbox_mask.png"
                    pil.save(out)
                    print(f"  person{j:02d}: score={float(score):.3f} cls={int(cls)} box=({x1},{y1},{x2},{y2}) -> {out}")
                    persons += 1
        except Exception as e:
            print(f"warning: failed to process result batch {i}: {e}")
            continue

    print(f"done: persons_detected={persons} masks_dir={save_root}")
    return 0


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