import cv2
import numpy as np
from ultralytics import YOLO
from typing import List, Tuple, Dict, Any

# Load YOLOv8 model for pothole and crack detection
model = YOLO("models/yolov8n.pt")

def detect_potholes_and_cracks(frame: np.ndarray) -> Tuple[List[Dict[str, Any]], np.ndarray]:
    """
    Detect potholes and cracks in the frame using YOLOv8.
    Args:
        frame: Input frame as a numpy array.
    Returns:
        Tuple of (list of detections, annotated frame).
    """
    # Perform inference
    results = model(frame, classes=[0, 1], conf=0.5)  # Class 0: pothole, Class 1: crack
    
    detections = []
    for i, r in enumerate(results[0].boxes):
        x_min, y_min, x_max, y_max = map(int, r.xyxy[0])
        conf = float(r.conf)
        cls = int(r.cls)
        dtype = "hole" if cls == 0 else "crack"
        label = f"{dtype.capitalize()} {i+1}"
        
        # Determine severity based on confidence and size
        area = (x_max - x_min) * (y_max - y_min)
        severity = "Severe" if area > 1000 or conf > 0.8 else "Moderate" if area > 500 or conf > 0.6 else "Mild"
        
        detections.append({
            "box": [x_min, y_min, x_max, y_max],
            "label": label,
            "type": dtype,
            "confidence": conf,
            "severity": severity
        })
    
    return detections, frame