import cv2
import numpy as np
from ultralytics import YOLO
import os
import random

# Load YOLOv8m-seg model for crack detection
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
MODEL_PATH = os.path.join(BASE_DIR, "../models/yolov8m-seg.pt")
model = YOLO(MODEL_PATH)

import random

def detect_cracks(frame, model):
    """
    Detect cracks in a frame using YOLOv8.
    Args:
        frame: Input frame (numpy array)
        model: YOLO model
    Returns:
        list: List of detected cracks with type, label, coordinates, confidence, and severity
    """
    # Run YOLOv8 inference for cracks
    results = model(frame)
    detected_items = []
    line_counter = 1  # Initialize counter for numbered labels

    for r in results:
        for box in r.boxes:
            conf = float(box.conf[0])
            if conf < 0.5:
                continue
            cls = int(box.cls[0])
            label = model.names[cls]
            if label != "crack":  # Process only cracks
                continue
            xyxy = box.xyxy[0].cpu().numpy()
            x_min, y_min, x_max, y_max = map(int, xyxy)

            # Simulate severity for cracks
            severity = random.choice(["low", "medium", "high"])
            
            # Add numbered label
            detection_label = f"Line {line_counter} - Crack (Conf: {conf:.2f})"
            item = {
                "type": label,
                "label": detection_label,
                "confidence": conf,
                "coordinates": [x_min, y_min, x_max, y_max],
                "severity": severity
            }

            detected_items.append(item)
            line_counter += 1

    return detected_items

def detect_potholes(frame, model):
    """
    Detect potholes in a frame using YOLOv8.
    Args:
        frame: Input frame (numpy array)
        model: YOLO model
    Returns:
        list: List of detected potholes with type, label, coordinates, and confidence
    """
    # Run YOLOv8 inference for potholes
    results = model(frame)
    detected_items = []
    line_counter = 1  # Initialize counter for numbered labels

    for r in results:
        for box in r.boxes:
            conf = float(box.conf[0])
            if conf < 0.5:
                continue
            cls = int(box.cls[0])
            label = model.names[cls]
            if label != "pothole":  # Process only potholes
                continue
            xyxy = box.xyxy[0].cpu().numpy()
            x_min, y_min, x_max, y_max = map(int, xyxy)

            # Add numbered label
            detection_label = f"Line {line_counter} - Pothole (Conf: {conf:.2f})"
            item = {
                "type": label,
                "label": detection_label,
                "confidence": conf,
                "coordinates": [x_min, y_min, x_max, y_max]
            }

            detected_items.append(item)
            line_counter += 1

    return detected_items

def detect_objects(frame, model):
    """
    Detect objects in a frame using YOLOv8.
    Args:
        frame: Input frame (numpy array)
        model: YOLO model
    Returns:
        list: List of detected objects with type, label, coordinates, and confidence
    """
    # Run YOLOv8 inference for other objects
    results = model(frame)
    detected_items = []
    line_counter = 1  # Initialize counter for numbered labels

    for r in results:
        for box in r.boxes:
            conf = float(box.conf[0])
            if conf < 0.5:
                continue
            cls = int(box.cls[0])
            label = model.names[cls]
            if label != "object":  # Process only objects
                continue
            xyxy = box.xyxy[0].cpu().numpy()
            x_min, y_min, x_max, y_max = map(int, xyxy)

            # Add numbered label
            detection_label = f"Line {line_counter} - Object (Conf: {conf:.2f})"
            item = {
                "type": label,
                "label": detection_label,
                "confidence": conf,
                "coordinates": [x_min, y_min, x_max, y_max]
            }

            detected_items.append(item)
            line_counter += 1

    return detected_items

def detect_items_in_sequence(frame, model):
    """
    Run crack, pothole, and object detection sequentially.
    Args:
        frame: Input frame (numpy array)
        model: YOLO model
    Returns:
        list: List of detected items (crack, pothole, object)
    """
    detected_items = []

    # Detect cracks first
    detected_items.extend(detect_cracks(frame, model))

    # Detect potholes second
    detected_items.extend(detect_potholes(frame, model))

    # Detect objects third
    detected_items.extend(detect_objects(frame, model))

    return detected_items