from ultralytics import YOLO
import os
import cv2
from PIL import Image


class ObjectDetector:
    def __init__(self, model_path="models/yolov11m.pt", threshold=0.5):
        """Initialize the object detector with YOLO model"""
        try:
            self.model = YOLO(model_path)
            self.threshold = threshold
            print(f"[INFO] Model loaded successfully: {model_path}")
        except Exception as e:
            print(f"❌ Error loading model: {e}")
            raise

    def detect_objects_in_folder(
        self, input_folder="core/input_frames", output_folder="core/output_frames"
    ):
        """Detect objects in all images in the input folder and save annotated images"""
        os.makedirs(output_folder, exist_ok=True)

        # Check if input folder exists and has images
        if not os.path.exists(input_folder):
            print(f"❌ Input folder not found: {input_folder}")
            return

        image_files = [
            f
            for f in os.listdir(input_folder)
            if f.lower().endswith((".jpg", ".jpeg", ".png"))
        ]

        if not image_files:
            print(f"❌ No image files found in: {input_folder}")
            return

        print(f"[INFO] Processing {len(image_files)} images...")

        for filename in image_files:
            image_path = os.path.join(input_folder, filename)

            try:
                # Run detection
                results = self.model(image_path, conf=self.threshold)

                if len(results) == 0:
                    print(f"[NO DETECTIONS] {filename}")
                    continue

                # Get detections
                result = results[0]
                if result.boxes is None or len(result.boxes) == 0:
                    print(f"[NO DETECTIONS] {filename}")
                    continue

                # Load original image for annotation
                original_image = cv2.imread(image_path)
                if original_image is None:
                    print(f"❌ Could not load image: {filename}")
                    continue

                # Get bounding boxes and draw them
                boxes = result.boxes.data.cpu().numpy()
                detected_count = 0

                for box in boxes:
                    x1, y1, x2, y2, conf, cls = box

                    # Convert to integers for drawing
                    x1, y1, x2, y2 = map(int, [x1, y1, x2, y2])

                    # Draw bounding box
                    cv2.rectangle(original_image, (x1, y1), (x2, y2), (0, 255, 0), 2)

                    # Add label
                    label = f"Drone: {conf:.2f}"
                    cv2.putText(
                        original_image,
                        label,
                        (x1, y1 - 10),
                        cv2.FONT_HERSHEY_SIMPLEX,
                        0.5,
                        (255, 0, 0),
                        2,
                    )

                    detected_count += 1

                # Save annotated image
                output_path = os.path.join(output_folder, filename)
                cv2.imwrite(output_path, original_image)
                print(f"[SAVED] {output_path} ({detected_count} objects)")

            except Exception as e:
                print(f"❌ Error processing {filename}: {e}")
                continue