import gradio as gr
import requests
import os

# --- IMPORTANT: UPDATE THIS URL ---
# After deploying your backend_app.py, paste its public URL here.
BACKEND_URL = "https://am-om-backend.hf.space/track/"
def track_video_via_api(input_video_path):
    if input_video_path is None:
        return None, "Please upload a video first."
    
    status = "Sending video to the backend for processing..."
    yield status, None, None # Update status immediately

    try:
        with open(input_video_path, "rb") as video_file:
            files = {'video': (os.path.basename(input_video_path), video_file, 'video/mp4')}
            
            # Send the video to the backend API and stream the response
            with requests.post(BACKEND_URL, files=files, stream=True, timeout= 600 ) as response:
                response.raise_for_status()
                
                # Save the processed video that the backend sends back
                processed_video_path = "processed_video.mp4"
                with open(processed_video_path, "wb") as f:
                    for chunk in response.iter_content(chunk_size=8192):
                        f.write(chunk)
                
                # For this Gradio app, we will just show the video.
                # You could modify the backend to return JSON as well.
                status = "Processing complete!"
                yield status, processed_video_path, None

    except requests.exceptions.RequestException as e:
        status = f"API Error: {e}"
        yield status, None, None

    except requests.exceptions.ReadTimeout:
        gr.Error("The processing took too long (more than 10 minutes). Please try a shorter video.")
        return None

    

# --- GRADIO UI (from your original app.py) ---
with gr.Blocks() as demo:
    gr.Markdown("<h1 style='text-align:center; font-size:4em;'>Vehicle and Pedestrian Segmentation</h1>") 
    gr.Markdown("Note: Processing on the free tier can take several minutes. Please be patient after clicking 'Start Tracking'.")
    input_video = gr.Video(label="Upload Video")
    start_btn = gr.Button("Start Tracking")
    status_display = gr.HTML("")  # Initially empty
    output_video = gr.Video(label="Processed Video")
    output_json = gr.File(label="Download JSON Output") # Note: This is not hooked up in this version

    start_btn.click(
        fn=track_video_via_api,
        inputs=input_video,
        outputs=[status_display, output_video, output_json]
    )

if __name__ == "__main__":
    demo.launch(share = True)