# 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 existing_file(path):
    try:
        p = Path(path).expanduser().resolve()
        if p.exists() and p.is_file() and p.stat().st_size > 0:
            return p
    except Exception:
        pass
    return None


def resolve_model_path(root, requested):
    if requested:
        p = existing_file(requested)
        if p:
            return p, "yolo26_model"
        return Path(requested).expanduser(), "detector_missing"

    trained = root / "detector" / "model" / "best.pt"
    p = existing_file(trained)
    if p:
        return p, "yolo26_detector"

    return trained, "detector_missing"


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--root", required=True)
    parser.add_argument("--image", required=True)
    parser.add_argument("--model", default="")
    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, model_source = resolve_model_path(root, args.model)
    if not model_path.exists():
        print(json.dumps({
            "available": False,
            "source": model_source,
            "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"

    try:
        results = model.predict(
            source=str(image_path),
            conf=float(args.conf),
            imgsz=int(args.imgsz),
            verbose=False,
            device=device,
        )
    except Exception as e:
        print(json.dumps({
            "available": False,
            "source": "detector_predict_failed",
            "modelPath": str(model_path),
            "image": str(image_path),
            "error": repr(e),
            "boxes": [],
        }, ensure_ascii=False))
        return

    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": model_source,
        "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()