import gradio as gr
import torch
import numpy as np
from scenedetect import detect, ContentDetector
from moviepy.editor import VideoFileClip
from transformers import AutoProcessor, AutoModel
import os
import uuid
from concurrent.futures import ThreadPoolExecutor

# Global variables and setup
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
processor = AutoProcessor.from_pretrained("microsoft/xclip-base-patch32")
model = AutoModel.from_pretrained("microsoft/xclip-base-patch32").to(device)

# To store the preprocessed scenes globally
preprocessed_scenes_global = None

def save_uploaded_file(uploaded_file):
    upload_dir = "uploaded_videos"
    os.makedirs(upload_dir, exist_ok=True)
    file_path = os.path.join(upload_dir, f"{uuid.uuid4()}.mp4")
    with open(file_path, "wb") as f:
        f.write(uploaded_file)
    return file_path

def extract_scenes(video_path):
    scene_list = detect(video_path, ContentDetector(threshold=32))
    return [(scene[0].get_seconds(), scene[1].get_seconds()) for scene in scene_list]

def extract_frames(video, start_time, end_time, num_frames=8):
    clip = video.subclip(start_time, end_time)
    frame_times = np.linspace(0, clip.duration, num_frames)
    return [clip.get_frame(t) for t in frame_times]

def preprocess_frames(frames):
    return processor(videos=frames, return_tensors="pt")

def preprocess_video(video_path):
    scenes = extract_scenes(video_path)
    video = VideoFileClip(video_path)

    preprocessed_scenes = []
    for start_time, end_time in scenes:
        frames = extract_frames(video, start_time, end_time)
        preprocessed_frames = preprocess_frames(frames)
        preprocessed_scenes.append({
            'frames': preprocessed_frames,
            'start_time': start_time,
            'end_time': end_time,
            'video_path': video_path  # Store video path for later use
        })

    video.close()
    return preprocessed_scenes

def run_inference(text_prompt):
    global preprocessed_scenes_global
    if preprocessed_scenes_global is None:
        return "Please preprocess the video first."

    text_inputs = processor(text=[text_prompt, "Nothing", "RedBull athlete in a blue hoodie standing on a ski hill", "Walking around in a skate park", "Logo on a T-shirt", "standing around inside", "Person talking to the camera", "Laying down on a snowy hill"], return_tensors="pt", padding=True).to(device)

    scene_scores = []
    for scene in preprocessed_scenes_global:
        with torch.no_grad():
            inputs = {k: v.to(device) for k, v in scene['frames'].items()}
            outputs = model(**inputs, **text_inputs)
            logits_per_video = outputs.logits_per_video
            probs = logits_per_video.softmax(dim=1)
            scene_score = probs[0][0].item()

        scene_scores.append({
            'start_time': scene['start_time'],
            'end_time': scene['end_time'],
            'score': scene_score,
            'video_path': scene['video_path']
        })

    sorted_scenes = sorted(scene_scores, key=lambda x: x['score'], reverse=True)
    return sorted_scenes[:3]  # Return the top 3 scenes

def process_video(video_file):
    global preprocessed_scenes_global
    if video_file:
        video_path = save_uploaded_file(video_file)
        preprocessed_scenes_global = preprocess_video(video_path)
        return "Video preprocessed successfully. You can now enter a text prompt."
    else:
        return "Please upload a video first."

def save_video_clip(scene, output_path):
    video = VideoFileClip(scene['video_path']).subclip(scene['start_time'], scene['end_time'])
    video.write_videofile(output_path, codec='libx264', audio_codec='aac', threads=4, preset='ultrafast')
    video.close()
    return output_path

def display_results(prompt):
    scenes = run_inference(prompt)
    if isinstance(scenes, str):  # Error message returned
        return [scenes] + [None] * 3
    
    output_paths = [f"/tmp/{uuid.uuid4()}.mp4" for _ in scenes]
    
    with ThreadPoolExecutor() as executor:
        output_clips = list(executor.map(save_video_clip, scenes, output_paths))
    
    return [""] + output_clips

CSS = """
    #col-container {margin-left: auto; margin-right: auto;}
    a {text-decoration-line: underline; font-weight: 600;}
"""

TITLE = """
    <div style="text-align: center; max-width: 650px; margin: 0 auto;">
        <h1 style="font-weight: 900; margin-bottom: 7px;">
            Video Clip Search
        </h1>
        <p style="margin-bottom: 10px; font-size: 94%">
            Upload a video, preprocess it, then enter a prompt to find relevant clips.
        </p>
    </div>
"""

with gr.Blocks(css=CSS) as demo:
    with gr.Column(elem_id="col-container"):
        gr.HTML(TITLE)

        with gr.Tab("Step 1: Upload and Preprocess Video"):
            with gr.Row():
                video_file = gr.File(label="Upload Video File:", type="binary")
            preprocess_button = gr.Button("Preprocess Video")
            preprocess_output = gr.Textbox(label="", interactive=False)

        with gr.Tab("Step 2: Enter Text Prompt"):
            with gr.Row():
                prompt_input = gr.Textbox(label="Enter Prompt", placeholder="Describe the scene you are looking for...")
            submit_button = gr.Button("Run Inference")
            output_clips = [gr.Video(label=f"Clip {i+1}") for i in range(3)]

        preprocess_button.click(fn=process_video, inputs=[video_file], outputs=preprocess_output)
        submit_button.click(fn=display_results, inputs=[prompt_input], outputs=[preprocess_output] + output_clips)

    demo.launch(share=True)