#!/usr/bin/env python3

import json
import os
import tempfile
import time
from pathlib import Path

from ultralytics import YOLO


IMAGE_PATH = Path("/dev/shm/robot_frame.jpg")
OUTPUT_PATH = Path("/dev/shm/r1_detections.json")

MODEL = "yolo26n.pt"

CONFIDENCE = 0.35
IMAGE_SIZE = 640


def atomic_write_json(path: Path, data):
    with tempfile.NamedTemporaryFile(
        mode="w",
        dir=path.parent,
        delete=False,
        suffix=".tmp",
    ) as f:
        json.dump(
            data,
            f,
            ensure_ascii=False,
            indent=2,
        )

        tmp = f.name

    os.replace(tmp, path)


def main():
    print(f"[YOLO] Loading {MODEL}")

    model = YOLO(MODEL)

    print("[YOLO] Ready")
    print(f"[YOLO] Camera: {IMAGE_PATH}")
    print(f"[YOLO] Output: {OUTPUT_PATH}")

    last_mtime = None

    while True:
        try:
            if not IMAGE_PATH.exists():
                time.sleep(0.05)
                continue

            mtime = IMAGE_PATH.stat().st_mtime_ns

            if mtime == last_mtime:
                time.sleep(0.01)
                continue

            last_mtime = mtime

            results = model.predict(
                source=str(IMAGE_PATH),
                conf=CONFIDENCE,
                imgsz=IMAGE_SIZE,
                verbose=False,
            )

            if not results:
                continue

            result = results[0]

            height, width = result.orig_shape

            detections = []

            if result.boxes is not None:
                for box in result.boxes:
                    class_id = int(box.cls[0])
                    confidence = float(box.conf[0])

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

                    cx = (x1 + x2) / 2.0
                    cy = (y1 + y2) / 2.0

                    detections.append({
                        "class_id": class_id,
                        "class": result.names[class_id],
                        "confidence": confidence,

                        "box": {
                            "x1": x1,
                            "y1": y1,
                            "x2": x2,
                            "y2": y2,
                        },

                        "center": {
                            "x": cx,
                            "y": cy,
                        },

                        "normalized_center": {
                            "x": cx / width,
                            "y": cy / height,
                        },
                    })

            data = {
                "timestamp": time.time(),
                "width": width,
                "height": height,
                "detections": detections,
            }

            atomic_write_json(
                OUTPUT_PATH,
                data,
            )

            objects = ", ".join(
                f"{d['class']} {d['confidence']:.2f}"
                for d in detections
            )

            if objects:
                print(f"[YOLO] {objects}")
            else:
                print("[YOLO] no objects")

        except KeyboardInterrupt:
            print("\n[YOLO] stopped")
            break

        except Exception as e:
            print(f"[YOLO] error: {e}")
            time.sleep(0.25)


if __name__ == "__main__":
    main()