import os
import glob
import subprocess
import shutil
import gradio as gr
from ultralytics import YOLO
from huggingface_hub import hf_hub_download

model_path = hf_hub_download(
    repo_id="deki/yolo26n-coco8-finetune",
    filename="best.pt",
    local_dir="models" 
)

model = YOLO(model_path)

def detect_objects(input_file, confidence: float = 0.25, is_video: bool = False):
    """Handle both image and video"""
        
    if is_video:
            
        # run detection on the video -> draw annotations frame-by-frame -> save the annotated video to disk
        results = model.predict(
        source=input_file,
        conf=confidence,
        save=True,
        project="/tmp",
        name="output_video",
        exist_ok=True,
        verbose=True,
        device="cpu",
        stream=True
        )
        
        last_result = None
        
        for result in results:
            last_result = result
        
        if last_result is None:
            raise gr.Error("No results generated.")
        
        save_dir = last_result.save_dir
        
        print("Save dir:", save_dir)
        print("Files:", os.listdir(save_dir))
        
        video_files = (
            glob.glob(os.path.join(save_dir, "*.mp4")) +
            glob.glob(os.path.join(save_dir, "*.avi"))
        )
        
        print("Generated files:", video_files)
        
        if not video_files:
            raise gr.Error("No MP4 video generated.")
        
        latest_video = max(video_files, key=os.path.getmtime)
        return latest_video
        
    else:
        # Process image
        results = model(input_file, conf=confidence, verbose=False) # user uploads image -> YOLO predicts boxes/classes/confidence scores 
        annotated_image = results[0].plot()                         # draws boxes onto the image 
        return annotated_image # Gradio displays annotated image
        


# === Improved Gradio Interface ===
with gr.Blocks(title="YOLO26 Object Detection") as demo:
    gr.Markdown("# 🚀 YOLO26 Fine-Tuned Object Detection")
    gr.Markdown("**Images + Short Video Support**")

    with gr.Tab("📷 Image Detection"):
        with gr.Row():
            with gr.Column():
                image_input = gr.Image(type="pil", label="Upload Image")
                conf_image = gr.Slider(0.1, 0.95, value=0.25, step=0.05, label="Confidence Threshold")
                image_btn = gr.Button("Detect Objects", variant="primary")
            image_output = gr.Image(label="Detections")
        
        image_btn.click(
            fn=lambda img, conf: detect_objects(img, conf, False),
            inputs=[image_input, conf_image],
            outputs=image_output
        )

    with gr.Tab("🎥 Video Detection"):
        with gr.Row():
            with gr.Column():
                video_input = gr.Video(label="Upload Short Video (10-15 seconds recommended)")
                conf_video = gr.Slider(0.1, 0.95, value=0.25, step=0.05, label="Confidence Threshold")
                video_btn = gr.Button("Process Video", variant="primary")
            video_output = gr.Video(label="Processed Video with Detections")
        
        video_btn.click(
            fn=lambda vid, conf: detect_objects(vid, conf, True),
            inputs=[video_input, conf_video],
            outputs=video_output
        )

    gr.Markdown(
    """
    ### About this Project
    - Fine-tuned YOLO26n on COCO8 (50 epochs) on M1 Pro MacBook
    - Supports both images and short videos
    """
    )

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