import gradio as gr
import os
import numpy as np
import cv2
import traceback
import logging
from PIL import Image
from helmet_detect_alert import detect_and_alert, MODEL_PATHS
from ultralytics import YOLO

# Load YOLOv8 model
model_path = MODEL_PATHS["YOLOv8n"]
model = YOLO(model_path)

# ---------- Detection Functions ----------
logging.basicConfig(filename="error_log.txt", level=logging.ERROR)

def detect_image_fn(image, confidence):
    try:
        if image is None:
            raise ValueError("No image uploaded.")
        
        # Convert PIL image to NumPy
        image_np = np.array(image.convert("RGB"))
        
        # YOLO prediction
        results = model.predict(image_np, conf=confidence, verbose=False)
        if not results:
            raise RuntimeError("Model returned empty results.")

        # Annotate result
        annotated = results[0].plot()
        annotated_rgb = cv2.cvtColor(annotated, cv2.COLOR_BGR2RGB)
        return Image.fromarray(annotated_rgb)

    except Exception as e:
        # Log error to file
        logging.error("Image detection failed:\n%s", traceback.format_exc())
        print("Error during prediction:", e)
        return None  # Gradio will show "Error"

def detect_video_fn(video_file, confidence):
    if video_file is None:
        return None
    input_path = video_file
    output_path = os.path.join("output", os.path.basename(input_path).replace(".", "_pred.")) + "mp4"
    os.makedirs("output", exist_ok=True)
    detect_and_alert(input_path, output_path, model, confidence)
    return output_path

# ---------- Interfaces ----------
image_interface = gr.Interface(
    fn=detect_image_fn,
    inputs=[
        gr.Image(label="Upload Image"),
        gr.Slider(0.1, 1.0, value=0.3, step=0.05, label="Confidence")
    ],
    outputs=gr.Image(label="Predicted Output"),
    title="🖼 Helmet Detection from Image",
    description="Upload an image to detect heads not wearing helmets."
)

video_interface = gr.Interface(
    fn=detect_video_fn,
    inputs=[
        gr.Video(label="Upload Video"),
        gr.Slider(0.1, 1.0, value=0.3, step=0.05, label="Confidence")
    ],
    outputs=gr.Video(label="Predicted Output"),
    title="🎥 Helmet Detection from Video",
    description="Upload a video and get helmet detection alerts visually."
)

# ---------- Launch Tabbed App ----------
gr.TabbedInterface(
    [image_interface, video_interface],
    tab_names=["Image Detection", "Video Detection"]
).launch()