# backend\ml\predict_detector_model.py

import argparse
import json
from pathlib import Path

from PIL import Image
from ultralytics import YOLO
import torch


def clamp01(v):
    return max(0.0, min(1.0, float(v)))


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--root", required=True)
    parser.add_argument("--image", required=True)
    parser.add_argument("--conf", type=float, default=0.30)
    parser.add_argument("--imgsz", type=int, default=640)
    parser.add_argument("--debug-image", action="store_true")
    args = parser.parse_args()

    root = Path(args.root)
    image_path = Path(args.image)

    model_path = root / "detector" / "model" / "best.pt"
    if not model_path.exists():
        print(json.dumps({
            "available": False,
            "source": "detector_missing",
            "modelPath": str(model_path),
            "boxes": [],
        }, ensure_ascii=False))
        return

    img = Image.open(image_path).convert("RGB")
    img_w, img_h = img.size

    try:
        model = YOLO(str(model_path))
    except Exception as e:
        print(json.dumps({
            "available": False,
            "source": "detector_load_failed",
            "modelPath": str(model_path),
            "error": repr(e),
            "boxes": [],
        }, ensure_ascii=False))
        return

    device = 0 if torch.cuda.is_available() else "cpu"

    results = model.predict(
        source=str(image_path),
        conf=float(args.conf),
        imgsz=int(args.imgsz),
        verbose=False,
        device=device,
    )

    boxes = []
    model_names = {}

    raw_box_count = 0

    if results:
        r = results[0]
        model_names = {str(k): str(v) for k, v in (r.names or {}).items()}

        if r.boxes is not None:
            raw_box_count = len(r.boxes)

            for b in r.boxes:
                cls_id = int(b.cls[0].item())
                score = float(b.conf[0].item())
                label = str((r.names or {}).get(cls_id, cls_id))

                x1, y1, x2, y2 = [float(v) for v in b.xyxy[0].tolist()]

                x1 = max(0.0, min(float(img_w), x1))
                y1 = max(0.0, min(float(img_h), y1))
                x2 = max(0.0, min(float(img_w), x2))
                y2 = max(0.0, min(float(img_h), y2))

                if x2 <= x1 or y2 <= y1:
                    continue

                x = x1 / img_w
                y = y1 / img_h
                w = (x2 - x1) / img_w
                h = (y2 - y1) / img_h

                if w <= 0 or h <= 0:
                    continue

                boxes.append({
                    "label": label,
                    "score": score,
                    "x": x,
                    "y": y,
                    "w": w,
                    "h": h,
                })

        if args.debug_image:
            debug_dir = root / "detector" / "debug"
            debug_dir.mkdir(parents=True, exist_ok=True)
            out_path = debug_dir / f"{image_path.stem}_conf_{args.conf}.jpg"
            plotted = r.plot()
            Image.fromarray(plotted).save(out_path)

    print(json.dumps({
        "available": True,
        "source": "yolo_detector",
        "modelPath": str(model_path),
        "image": str(image_path),
        "conf": float(args.conf),
        "imgsz": int(args.imgsz),
        "device": str(device),
        "imageWidth": img_w,
        "imageHeight": img_h,
        "classNames": model_names,
        "rawBoxCount": raw_box_count,
        "boxes": boxes,
    }, ensure_ascii=False))


if __name__ == "__main__":
    main()