# backend\ml\predict_detector_model.py

import argparse
import json
from pathlib import Path

from PIL import Image
from ultralytics import YOLO


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)
    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",
            "boxes": [],
        }))
        return

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

    model = YOLO(str(model_path))
    results = model.predict(
        source=str(image_path),
        conf=args.conf,
        imgsz=args.imgsz,
        verbose=False,
        device=0 if __import__("torch").cuda.is_available() else "cpu",
    )

    boxes = []

    if results:
        r = results[0]
        names = r.names or {}

        if r.boxes is not None:
            for b in r.boxes:
                cls_id = int(b.cls[0].item())
                score = float(b.conf[0].item())
                label = str(names.get(cls_id, cls_id))

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

                x = clamp01(x1 / img_w)
                y = clamp01(y1 / img_h)
                w = clamp01((x2 - x1) / img_w)
                h = clamp01((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,
                })

    print(json.dumps({
        "available": True,
        "source": "yolo_detector",
        "boxes": boxes,
    }, ensure_ascii=False))


if __name__ == "__main__":
    main()