from ultralytics import YOLO
from PIL import Image, ImageDraw, ImageFont
import numpy as np


class RoadsideDetector:
    def __init__(self, model_path="yol_kenari_tespit_best.pt"):
        self.model = YOLO(model_path)

        # Türkçe sınıf isimleri
        self.class_names = {
            0: "çöp",
            1: "çöp kutusu",
            2: "taşan çöp",
        }

        # Seviye hesaplamasına giren sınıflar
        self.severity_classes = ["çöp", "taşan çöp"]

    def detect_items(self, image):
        results = self.model(image, verbose=False)
        detections = []

        for r in results:
            boxes = r.boxes.xyxy.cpu().numpy()
            confs = r.boxes.conf.cpu().numpy()
            classes = r.boxes.cls.cpu().numpy()

            for box, conf, cls_id in zip(boxes, confs, classes):
                cls_name = self.class_names.get(int(cls_id), "Bilinmeyen")
                detections.append({
                    "class_name": cls_name,
                    "conf": float(conf),
                    "box": box.tolist(),
                })
        return detections

    def draw_boxes(self, image, detections):
        """Türkçe karakter desteği ile bounding box çizimi"""
        if not isinstance(image, Image.Image):
            image = Image.fromarray(image)

        draw = ImageDraw.Draw(image)
        try:
            font = ImageFont.truetype("DejaVuSans.ttf", 28)
        except:
            font = ImageFont.load_default()

        for det in detections:
            box = det["box"]
            cls_name = det["class_name"]
            conf = det["conf"]

            x1, y1, x2, y2 = box
            draw.rectangle([x1, y1, x2, y2], outline="lime", width=4)

            label = f"{cls_name} ({conf:.2f})"

            # ✅ Yeni Pillow sürümü için textsize yerine textbbox kullanalım
            bbox = draw.textbbox((0, 0), label, font=font)
            text_w = bbox[2] - bbox[0]
            text_h = bbox[3] - bbox[1]

            draw.rectangle([x1, y1 - text_h - 4, x1 + text_w + 4, y1], fill="black")
            draw.text((x1 + 2, y1 - text_h - 2), label, font=font, fill="yellow")

        return image

    def calculate_severity(self, detections):
        """Sadece çöp ve taşan çöp seviye hesabına girer"""
        severity_count = sum(1 for d in detections if d["class_name"] in self.severity_classes)
        return severity_count