import spaces
import gradio as gr
import shutil
import os
import subprocess
import sys
# Run the .bat file before launching the app
try:
    import PromptTrack
except ImportError:
    print("PromptTrack not found. Installing...")
    subprocess.run([sys.executable, "-m", "pip", "install", 
                    "--index-url", "https://test.pypi.org/simple/", 
                    "--extra-index-url", "https://pypi.org/simple/", 
                    "PromptTrack"], check=True)
    subprocess.run([sys.executable, "-m", "pip", "install", 
                    "--no-deps", "bytetracker"], check=True)
    import PromptTrack  # Retry import after installation


from PromptTrack import PromptTracker
tracker = PromptTracker()


#@spaces.GPU(duration=300)
def process_video(video_path, prompt):
    import torch
    print(f"Is CUDA available: {torch.cuda.is_available()}")
    # True
    print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
    # Tesla T4
    detection_threshold=0.3
    track_thresh=0.4
    match_thresh=1
    max_time_lost=float("inf")
    nbr_frames_fixing=800
    output_video = video_path.split('mp4')[0]+"_with_id.mp4"  # Placeholder for processed video
    output_file = video_path.split('mp4')[0]+"_mot_.json"    # Tracking result
    output_file_2 = video_path.split('mp4')[0]+"_object_detection.json"    # detection results
    video_file = video_path
    tracker.detect_objects(video_file, prompt=prompt, nms_threshold=0.8, detection_threshold=detection_threshold, detector="OWL-VITV2")
    tracker.process_mot(video_file, fixed_parc=True, track_thresh=track_thresh, match_thresh=match_thresh, frame_rate=25, max_time_lost=max_time_lost, nbr_frames_fixing=nbr_frames_fixing)
    tracker.read_video_with_mot(video_file, fps=25)
    
    
    """output_video = "output.mp4"  # Placeholder for processed video
    output_file = "output.txt"    # Placeholder for generated file
    
    
    # Copy the input video to simulate processing
    shutil.copy(video_path.name, output_video)
    
    # Create an output text file with the prompt content
    with open(output_file, "w") as f:
        f.write(f"User Prompt: {prompt}\n")
    """
    return output_video, output_file

# Define Gradio interface
iface = gr.Interface(
    fn=process_video,
    inputs=[gr.File(label="Upload Video"), gr.Textbox(placeholder="Enter your prompt")],
    outputs=[gr.Video(), gr.File(label="Generated File")],
    title="Video Processing App",
    description="Upload a video and enter a prompt. The app will return the processed video and a generated file."
)


# Launch the app
if __name__ == "__main__":
    iface.launch(share=True)