from __future__ import annotations

from pathlib import Path
from time import perf_counter

import cv2
import gradio as gr
from ultralytics import YOLO


MODEL_PATH = "yolov8n.pt" if Path("yolov8n.pt").exists() else "yolo11n.pt"

# Initial COCO proxy classes:
# 29 = frisbee, 39 = bottle, 54 = donut
# These are not biscuit classes. They are only proxies for round biscuit-like shapes.
BISCUIT_LIKE_CLASS_IDS = {29, 39, 54}
INFERENCE_IMAGE_SIZE = 320

model = YOLO(MODEL_PATH)


def detect_biscuit_defects(image, confidence):
    started_at = perf_counter()

    if image is None or image.size == 0:
        return None, "Waiting for webcam frame."

    bgr_image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    results = model(
        bgr_image,
        conf=confidence,
        imgsz=INFERENCE_IMAGE_SIZE,
        max_det=10,
        verbose=False,
    )[0]

    detected_pieces = []
    for box in results.boxes:
        class_id = int(box.cls[0])
        if class_id in BISCUIT_LIKE_CLASS_IDS:
            detected_pieces.append(box)

    if len(detected_pieces) == 1:
        box = detected_pieces[0]
        x1, y1, x2, y2 = map(int, box.xyxy[0])
        draw_box(bgr_image, x1, y1, x2, y2, "INTACT BISCUIT", float(box.conf[0]), (0, 255, 0))
        status = "Proxy method: 1 biscuit-like object -> intact"
    elif len(detected_pieces) >= 2:
        for box in detected_pieces:
            x1, y1, x2, y2 = map(int, box.xyxy[0])
            draw_box(bgr_image, x1, y1, x2, y2, "BROKEN PIECE", float(box.conf[0]), (0, 0, 255))
        status = f"Proxy method: {len(detected_pieces)} biscuit-like objects -> broken"
    else:
        status = "Proxy method: no biscuit-like object detected"

    elapsed_ms = int((perf_counter() - started_at) * 1000)
    status = f"{status} | {elapsed_ms} ms/frame at imgsz={INFERENCE_IMAGE_SIZE}"

    return cv2.cvtColor(bgr_image, cv2.COLOR_BGR2RGB), status


def draw_box(bgr_image, x1, y1, x2, y2, label, confidence, color):
    cv2.rectangle(bgr_image, (x1, y1), (x2, y2), color, 3)
    cv2.putText(
        bgr_image,
        f"{label} {confidence:.2f}",
        (x1, max(y1 - 10, 25)),
        cv2.FONT_HERSHEY_SIMPLEX,
        0.8,
        color,
        2,
        cv2.LINE_AA,
    )


with gr.Blocks(title="YOLO Proxy Biscuit Detector") as demo:
    gr.Markdown("# YOLO Proxy Biscuit Detector")
    gr.Markdown(
        f"Loaded proxy model: `{MODEL_PATH}`. "
        "This uses COCO proxy classes, not custom biscuit weights."
    )

    with gr.Row():
        webcam = gr.Image(
            sources=["webcam"],
            type="numpy",
            streaming=True,
            label="Webcam",
        )
        annotated = gr.Image(type="numpy", label="Proxy detection")

    confidence = gr.Slider(
        0.05,
        0.80,
        value=0.20,
        step=0.05,
        label="Confidence threshold",
    )
    status = gr.Textbox(label="Status")

    webcam.stream(
        fn=detect_biscuit_defects,
        inputs=[webcam, confidence],
        outputs=[annotated, status],
        stream_every=0.10,
        queue=False,
    )


if __name__ == "__main__":
    demo.launch()