import cv2
import gradio as gr
import uuid
import os
import torch
import tempfile
import shutil
from ultralytics import YOLO

# Automatically download your best.pt model from your dataset repo
model_path = "best.pt"
if not os.path.exists(model_path):
    os.system("wget https://huggingface.co/datasets/Prasanna1622/solar-fault-dataset/resolve/main/best.pt")

# Initialize the YOLO model
model = YOLO(model_path)

# Inference function
def detect_faults(video_path):
    """
    - video_path: the path to the uploaded video file on disk.
    - Returns: path to the annotated output.mp4.
    """
    try:
        # Create a unique RUN directory so YOLO does not overwrite previous results
        unique_id = str(uuid.uuid4())[:8]
        project_dir = os.path.join("runs", "detect", unique_id)
        os.makedirs(project_dir, exist_ok=True)
        print(f"🛠️  Running inference, saving to: {project_dir}")

        # Run YOLO predict; this saves the annotated video in project_dir/
        results = model.predict(
            source=video_path,  # path to the uploaded video
            save=True,
            save_txt=False,
            conf=0.5,
            project=os.path.join("runs", "detect"),
            name=unique_id
        )
        print("✅ YOLO predict() finished.")

        # Check if output video exists
        original_name = os.path.basename(video_path)
        output_video_path = os.path.join("runs", "detect", unique_id, original_name)
        print(f"🛠️  Looking for output video at: {output_video_path}")

        if os.path.exists(output_video_path):
            print("✅ Output video found, returning it.")
            return output_video_path
        else:
            print(f"❌ Output video NOT found at: {output_video_path}")
            return "Error: Annotated video not found."

    except Exception as e:
        # Print the full exception in logs, return a simple string in UI
        print(f"❌ Exception during detect_faults: {e}")
        return "Error during processing."


# Create Gradio UI
demo = gr.Interface(
    fn=detect_faults,
    inputs=gr.Video(label="Upload Input Video"),
    outputs=gr.Video(label="Detected Output Video"),
    title="Solar Panel Fault Detection from Drone Video",
    description="Upload a drone video to detect solar panel faults using a YOLOv8 model."
)

if __name__ == "__main__":
    demo.launch()