import gradio as gr
import os
import shutil
from tracker import vehicle_tracker_and_counter
import tempfile

def process_video(video_path):
    if not video_path:
        return None
    
    # Create temp file for output
    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp_out:
        output_path = tmp_out.name
    
    try:
        # Run tracker
        # Initialize tracker with input and output paths
        tracker = vehicle_tracker_and_counter(
            source_video_path=video_path,
            target_video_path=output_path,
            use_tensorrt=False
        )
        
        tracker.run()
        return output_path
    except Exception as e:
        return f"Error: {str(e)}"

# Define CSS for centering and styling
css = """
#col-container {
    margin: 0 auto;
    max-width: 960px;
}
"""

with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown("# 🚗 SpeedLens: Vehicle Tracking & Speed Estimation")
        gr.Markdown("Upload a video to detect vehicles, track them, and estimate their speed.")
        
        with gr.Row():
            input_video = gr.Video(label="Input Video", sources=["upload"])
            output_video = gr.Video(label="Processed Video")
        
        process_btn = gr.Button("🚀 Process Video", variant="primary")
        
        process_btn.click(
            fn=process_video,
            inputs=input_video,
            outputs=output_video
        )
        
        gr.Markdown("### How it works")
        gr.Markdown("- **Detection**: Uses YOLOv8x to detect vehicles.")
        gr.Markdown("- **Tracking**: Uses ByteTrack for multi-object tracking.")
        gr.Markdown("- **Speed Estimation**: Uses perspective transformation based on calibrated road points.")

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)