"""
Open-vocabulary object detection via YOLO-World (YOLOv8 family, ultralytics),
run alongside SAM (see pipeline/flower_count.py) for the flower-count feature.

SAM has no notion of "flower" -- it blindly proposes every objectlike region
it can find, so the vase and any ribbon/bow tied around it get segmented and
counted as if they were flowers. YOLO-World, prompted with plain-language
classes, actually knows what those things look like: its "vase"/"ribbon"/
"bow" detections are used (in engine.count_flowers) to exclude those regions
from SAM's flower-instance count.

Its own "flower" boxes are shown to the user as a second, independent count
-- in practice it tends to draw one box per contiguous flower region rather
than per individual bloom, so it's a coarser, corroborating signal, not a
replacement for SAM's per-instance count.
"""

import gc
import logging

import cv2
import torch

import config

logger = logging.getLogger(__name__)

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

_model = None


def get_model():
    global _model
    if _model is None:
        logger.info("Loading YOLO-World (%s) on %s...", config.YOLO_WORLD_CHECKPOINT, DEVICE)
        from ultralytics import YOLO
        _model = YOLO(config.YOLO_WORLD_CHECKPOINT)
        _model.set_classes(config.YOLO_WORLD_CLASSES)
    return _model


def unload_model():
    global _model
    freed = _model is not None
    _model = None
    if freed:
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
    return freed


def detect(bgr):
    """Returns {class_name: [{"xyxy": [x1,y1,x2,y2], "confidence": float}, ...]}
    for every class in config.YOLO_WORLD_CLASSES (empty list if none found)."""
    model = get_model()
    with torch.no_grad():
        results = model.predict(bgr, conf=config.YOLO_WORLD_CONF,
                                 iou=config.YOLO_WORLD_IOU, verbose=False)
    r = results[0]
    by_class = {name: [] for name in config.YOLO_WORLD_CLASSES}
    for b in r.boxes:
        name = r.names[int(b.cls[0])]
        by_class.setdefault(name, []).append({
            "xyxy": [float(v) for v in b.xyxy[0].tolist()],
            "confidence": round(float(b.conf[0]), 3),
        })
    return by_class


_BOX_COLORS = {
    "flower": (66, 133, 244), "vase": (219, 68, 55),
    "ribbon": (244, 180, 0), "bow": (15, 157, 88),
}


def render_boxes(bgr, by_class):
    """Draws a labeled box per detection, colored by class, so the user can
    see exactly which regions YOLO-World identified as flower vs. vase vs.
    ribbon/bow."""
    overlay = bgr.copy()
    for name, boxes in by_class.items():
        color = _BOX_COLORS.get(name, (170, 170, 170))
        for box in boxes:
            x1, y1, x2, y2 = [int(v) for v in box["xyxy"]]
            cv2.rectangle(overlay, (x1, y1), (x2, y2), color, 3)
            label = f"{name} {box['confidence']:.2f}"
            (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.7, 2)
            label_y = max(th + 6, y1)
            cv2.rectangle(overlay, (x1, label_y - th - 8), (x1 + tw + 6, label_y), color, -1)
            cv2.putText(overlay, label, (x1 + 3, label_y - 5),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2, cv2.LINE_AA)
    return overlay