import gradio as gr
import cv2
import numpy as np
from ultralytics import YOLO
import os
from PIL import Image

# Load model
model_path = "models/aov_herodetector_v1.pt"
try:
    model = YOLO(model_path)
    print(f"Model loaded successfully from {model_path}")
except Exception as e:
    print(f"Error loading model: {e}")
    model = None

def detect_heroes(image):
    """
    Detect AOV heroes in the uploaded image
    """
    if model is None:
        return None, "Model not loaded properly"
    
    try:
        # Convert PIL Image to numpy array if needed
        if isinstance(image, Image.Image):
            image = np.array(image)
        
        # Run inference
        results = model(image)
        
        # Get the first result
        result = results[0]
        
        # Plot results on image
        annotated_image = result.plot()
        
        # Convert BGR to RGB (OpenCV uses BGR, but gradio expects RGB)
        annotated_image = cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)
        
        # Get detection info
        detection_info = ""
        if len(result.boxes) > 0:
            detection_info = f"Detected {len(result.boxes)} heroes:\n"
            for i, box in enumerate(result.boxes):
                conf = box.conf[0].item()
                cls = int(box.cls[0].item())
                class_name = model.names[cls] if cls < len(model.names) else f"Class_{cls}"
                detection_info += f"- {class_name}: {conf:.2f}\n"
        else:
            detection_info = "No heroes detected"
        
        return annotated_image, detection_info
        
    except Exception as e:
        return None, f"Error during detection: {str(e)}"

# Create Gradio interface
with gr.Blocks(title="AOV Hero Detector", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🎮 Arena of Valor Hero Detector")
    gr.Markdown("Upload an image to detect AOV heroes using YOLOv8 model")
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(
                label="Upload Image",
                type="pil",
                height=400
            )
            detect_btn = gr.Button("🔍 Detect Heroes", variant="primary")
            
        with gr.Column():
            output_image = gr.Image(
                label="Detection Result",
                height=400
            )
            detection_text = gr.Textbox(
                label="Detection Info",
                lines=5,
                max_lines=10
            )
    
    # Examples (optional - you can add sample images)
    gr.Examples(
        examples=[],  # Add paths to example images if you have them
        inputs=input_image,
        label="Example Images"
    )
    
    # Event handlers
    detect_btn.click(
        fn=detect_heroes,
        inputs=input_image,
        outputs=[output_image, detection_text]
    )
    
    # Auto-detect when image is uploaded
    input_image.change(
        fn=detect_heroes,
        inputs=input_image,
        outputs=[output_image, detection_text]
    )

# Launch the app
if __name__ == "__main__":
    demo.launch()