import os
import glob
import gradio as gr
import shutil
from ultralytics import YOLO

# Initialize YOLO model
model = YOLO('./best.pt')

def check_ffmpeg():
    """Check if ffmpeg is installed."""
    if shutil.which("ffmpeg") is None:
        raise EnvironmentError("ffmpeg is not installed or not found in the system PATH. Please install ffmpeg to proceed.")

def process_video(video_path, output_option):
    # Check if ffmpeg is installed
    check_ffmpeg()

    # Run the YOLO model on the video, specifying the tracker configuration
    results = model.track(video_path, save=True, tracker="bytetrack.yaml", persist=True)

    # Get the video file name without extension and directory
    video_name = os.path.splitext(os.path.basename(video_path))[0]
    
    # Find the latest directory created in 'runs/detect/'
    output_dir = max(glob.glob('./runs/detect/*'), key=os.path.getmtime)
    
    # Find the saved video file in the latest directory with the specific video name
    video_files = glob.glob(os.path.join(output_dir, f'{video_name}*.avi'))
    
    if not video_files:
        raise Exception(f"No .avi video files found in directory {output_dir} for {video_name}")

    # Convert the video to mp4
    mp4_files = []
    for file in video_files:
        mp4_file = f"{file[:-4]}.mp4"
        # Quote the file paths to handle spaces and overwrite without prompt
        os.system(f"ffmpeg -y -i \"{file}\" -vcodec libx264 \"{mp4_file}\"")
        os.remove(file)  # Remove the original .avi file after conversion
        mp4_files.append(mp4_file)

    # Check again for the converted .mp4 video files specifically for the processed video
    matched_mp4_files = [file for file in mp4_files if video_name in file]
    
    if not matched_mp4_files:
        raise Exception(f"No .mp4 video files found in directory {output_dir} after conversion for {video_name}.")

    # Initialize object count
    object_count = 0

    # Calculate object count based on results
    for result in results:
        if hasattr(result, 'boxes'):
            object_count += len(result.boxes)

    # Set default outputs
    video_output = matched_mp4_files[0] if matched_mp4_files else None
    count_output = object_count

    # Determine the output based on user's choice
    if output_option == "Count":
        return count_output, None
    elif output_option == "Video":
        return None, video_output
    elif output_option == "Both":
        return count_output, video_output

# Define Gradio inputs
video_input = gr.Video()
output_option_input = gr.Radio(choices=["Count", "Video", "Both"], label="Select output type")

# Define examples for the interface
examples = [
    [os.path.abspath("example_video1.mp4"), "Count"],
    [os.path.abspath("example_video2.mp4"), "Video"],
    [os.path.abspath("example_video3.mp4"), "Both"]
]

video_interface = gr.Interface(
    fn=process_video,
    inputs=[video_input, output_option_input],
    outputs=[gr.Textbox(label="Object Count"), "video"],  # Ensure two outputs are defined
    title="YOLO Video Tracking Application",
    description="A simple application to track objects in a video using YOLO model. Upload your own video, or click one of the examples to load them.",
    article="""<div>
                <p style="text-align: center">Upload a video file and select the type of output you want: object count, processed video, or both. Then, hit submit to process the video.</p>
               </div>""",
    examples=examples
)

# Deploy the interface with share enabled
gr.TabbedInterface([video_interface], ["Track Video"]).launch(share=True)