import gradio as gr
import cv2
import numpy as np
import tempfile
from ultralytics import YOLO

# --- Load YOLO Models ---
helmet_model = YOLO("models/helmet_detection.pt")
fire_model = YOLO("models/fire_detection.pt")
seg_model = YOLO("models/yolo11x-seg.pt")

model_map = {
    "Helmet Detection": helmet_model,
    "Fire Detection": fire_model,
    "Instance Segmentation": seg_model,
    "Workout Pose (MediaPipe)": None  # MediaPipe handled separately
}

# --- Placeholder MediaPipe workout processing ---
def process_workout(reference_video_path, user_video_path):
    # Here you would add your actual MediaPipe logic
    # For demo, just return the user video as "processed"
    return user_video_path

# --- Video processing function ---
def process_video(model_name, video_file, reference_file=None):
    if not video_file:
        return None
    
    # MediaPipe workflow
    if model_name == "Workout Pose (MediaPipe)":
        if not reference_file:
            return None
        return process_workout(reference_file.name, video_file.name)
    
    # Use the uploaded video path directly
    video_path = video_file.name
    model = model_map[model_name]
    
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        return None
    
    output_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4").name
    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    out = None
    
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        
        # Run YOLO prediction
        results = model.predict(frame)
        for r in results:
            if hasattr(r, "boxes"):
                boxes = r.boxes.xyxy.cpu().numpy()
                for box in boxes:
                    x1, y1, x2, y2 = map(int, box)
                    cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
        
        if out is None:
            h, w, _ = frame.shape
            out = cv2.VideoWriter(output_file, fourcc, 20.0, (w, h))
        out.write(frame)
    
    cap.release()
    if out:
        out.release()
    
    return output_file

# --- Gradio Interface ---
with gr.Blocks() as demo:
    gr.Markdown("# Multi-Model AI Video Analyzer")
    
    with gr.Row():
        model_choice = gr.Dropdown(list(model_map.keys()), label="Select Model")
        video_input = gr.File(label="Upload Video")
        reference_input = gr.File(label="Reference Video (for workout)", visible=False)
    
    output_video = gr.Video(label="Processed Video")
    
    def toggle_reference_input(model_name):
        return gr.update(visible=(model_name == "Workout Pose (MediaPipe)"))
    
    model_choice.change(toggle_reference_input, inputs=model_choice, outputs=reference_input)
    
    submit_btn = gr.Button("Process Video")
    submit_btn.click(
        process_video,
        inputs=[model_choice, video_input, reference_input],
        outputs=output_video
    )

demo.launch()