"""
Quick license plate detection smoke test.

Usage:
  python -m tools.test_license_plates <image_path>

Tests the license plate detection using the YOLO model at:
  models/license-plate-detection/lp_smallobj_1280/weights/best.pt

Environment variables:
  YOLO_LP_MODEL_PATH (optional absolute/relative path to .pt)
  YOLO_LP_MODEL (default: models/license-plate-detection/lp_smallobj_1280/weights/best.pt)
  YOLO_DEVICE (e.g., "cpu" or "0")
  YOLO_LP_CONF (default: 0.25)

Outputs:
  - Prints detection summary to stdout
  - Saves annotated image to artifacts/tmp/lp_debug/<image_stem>_annotated.jpg
  - Saves individual license plate masks as PNGs
"""
from __future__ import annotations

import os
import sys
from pathlib import Path
import argparse


def main() -> int:
    ap = argparse.ArgumentParser(description="Test license plate detection with YOLO")
    ap.add_argument("image", help="Path to input image")
    ap.add_argument("--conf", type=float, help="Confidence threshold (default: 0.10 - maximum sensitivity)")
    ap.add_argument("--imgsz", type=int, help="Inference image size (default: 2560 - maximum sensitivity for small plates)")
    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
        import numpy as np  # type: ignore
    except Exception as e:
        print(f"error: PIL/numpy not available: {e}")
        return 3
    
    try:
        import torch  # type: ignore
        from ultralytics import YOLO  # type: ignore
    except Exception as e:
        print(f"error: ultralytics import failed: {e}")
        return 4

    # Set up paths
    repo_root = Path(__file__).resolve().parents[1]
    default_model_path = repo_root / "models" / "license-plate-detection" / "lp_smallobj_1280" / "weights" / "best.pt"
    
    model_path = os.getenv("YOLO_LP_MODEL_PATH") or os.getenv("YOLO_LP_MODEL") or str(default_model_path)
    device = os.getenv("YOLO_DEVICE", "cpu")
    conf = args.conf if args.conf is not None else float(os.getenv("YOLO_LP_CONF", "0.10"))
    imgsz = args.imgsz if args.imgsz is not None else int(os.getenv("YOLO_LP_IMGSZ", "2560"))

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

    # Check if model exists
    if not Path(model_path).exists():
        print(f"error: model not found at {model_path}")
        print(f"Please ensure the license plate detection model is available")
        return 5

    try:
        model = YOLO(model_path)
        print("loaded license plate detection model")
    except Exception as e:
        print(f"license plate model load failed: {e}")
        return 6

    # Run inference
    try:
        print(f"running inference on {ipath}...")
        results = model.predict(
            source=str(ipath),
            conf=conf,
            imgsz=imgsz,
            device=device,
            verbose=False,
        )
    except Exception as e:
        print(f"inference failed: {e}")
        return 7

    if not results or len(results) == 0:
        print("no results returned")
        return 8

    result = results[0]
    boxes = result.boxes
    
    if boxes is None or len(boxes) == 0:
        print("no license plates detected")
        return 0

    print(f"\nDetected {len(boxes)} license plate(s):")
    
    # Load original image
    img = Image.open(ipath)
    img_draw = img.copy()
    draw = ImageDraw.Draw(img_draw)
    
    # Create output directory
    out_dir = Path("artifacts") / "tmp" / "lp_debug" / ipath.stem
    out_dir.mkdir(parents=True, exist_ok=True)
    
    # Process each detection
    for i, box in enumerate(boxes):
        xyxy = box.xyxy[0].cpu().numpy()
        conf_score = float(box.conf[0].cpu().numpy())
        
        x1, y1, x2, y2 = map(int, xyxy)
        w = x2 - x1
        h = y2 - y1
        
        print(f"  [{i}] bbox=({x1},{y1},{w},{h}) conf={conf_score:.3f}")
        
        # Draw bounding box on image
        draw.rectangle([x1, y1, x2, y2], outline="red", width=3)
        draw.text((x1, y1 - 10), f"LP {i} ({conf_score:.2f})", fill="red")
        
        # Save individual mask
        mask = Image.new("L", img.size, 0)
        mask_draw = ImageDraw.Draw(mask)
        mask_draw.rectangle([x1, y1, x2, y2], fill=255)
        
        mask_path = out_dir / f"license_plate_{i}_mask.png"
        mask.save(mask_path)
        print(f"      mask saved: {mask_path}")
    
    # Save annotated image
    annotated_path = out_dir / f"{ipath.stem}_annotated.jpg"
    img_draw.save(annotated_path)
    print(f"\nannotated image saved: {annotated_path}")
    
    return 0


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