import gradio as gr
import requests
import time
import json

# Configuration - NOTE the /run endpoint
RUNPOD_API_URL = "https://api.runpod.ai/v2/vz476hgvqvzojy/run"
RUNPOD_API_KEY = "rpa_0PN9A3W28RFTK3LDY6M3LRO5HT9SWI441HQV2V6A1jt5ih"

def process_video(video_url):
    """
    Step 1: Submit video URL to RunPod API
    """
    headers = {
        "Authorization": f"Bearer {RUNPOD_API_KEY}",
        "Content-Type": "application/json"
    }
    
    payload = {
        "input": {
            "video_url": video_url
        }
    }
    
    try:
        # Initial request to start processing
        yield "🔄 Submitting job to RunPod API..."
        response = requests.post(RUNPOD_API_URL, json=payload, headers=headers)
        response.raise_for_status()
        initial_response = response.json()
        
        job_id = initial_response.get('id')
        status = initial_response.get('status')
        
        if not job_id:
            yield f"❌ Error: No job ID received. Response: {initial_response}"
            return
        
        yield f"✅ Job submitted!\nID: {job_id}\nStatus: {status}\n\n🔄 Starting status polling..."
        
        # Poll for status updates
        max_attempts = 200  # Increased for longer processing times
        for attempt in range(max_attempts):
            time.sleep(5)  # Wait 5 seconds between checks
            
            # Status URL - NOTE: Remove /run from the base URL for status checks
            status_url = f"https://api.runpod.ai/v2/vz476hgvqvzojy/status/{job_id}"
            status_response = requests.get(status_url, headers=headers)
            
            if status_response.status_code != 200:
                yield f"❌ Status check failed (Ping {attempt+1}): HTTP {status_response.status_code}"
                continue
                
            status_data = status_response.json()
            current_status = status_data.get('status')
            
            if current_status == 'IN_QUEUE':
                yield f"⏳ Attempt {attempt+1}/{max_attempts}\nStatus: IN_QUEUE\nJob ID: {job_id}\n\nJob is waiting in queue..."
            elif current_status == 'IN_PROGRESS':
                delay_time = status_data.get('delayTime', 'N/A')
                yield f"🔄 Attempt {attempt+1}/{max_attempts}\nStatus: IN_PROGRESS\nJob ID: {job_id}\nDelay Time: {delay_time}\n\nProcessing in progress..."
            elif current_status == 'COMPLETED':
                output = status_data.get('output', {})
                yield f"✅ PROCESSING COMPLETE!\n\nFinal Output:\n{json.dumps(output, indent=2)}"
                return
            elif current_status == 'FAILED':
                error_msg = status_data.get('error', 'No error message')
                yield f"❌ JOB FAILED\nError: {error_msg}"
                return
            else:
                yield f"❓ Unknown status: {current_status}\nFull response:\n{json.dumps(status_data, indent=2)}"
                return
                
        yield f"⏰ Maximum polling attempts ({max_attempts}) reached. Job may still be processing."
            
    except requests.exceptions.RequestException as e:
        yield f"❌ Network error: {str(e)}"
    except Exception as e:
        yield f"❌ Unexpected error: {str(e)}"

# Create Gradio interface
with gr.Blocks(title="Video Processing API") as demo:
    gr.Markdown("""
    # Video Processing Interface
    Submit a Google Drive video URL for processing. The system will automatically track the job status and display the final results.
    """)
    
    with gr.Row():
        with gr.Column():
            video_url = gr.Textbox(
                label="Google Drive Video URL",
                placeholder="https://drive.google.com/file/d/...",
                lines=1
            )
            submit_btn = gr.Button("Process Video", variant="primary")
        
        with gr.Column():
            output = gr.Textbox(
                label="Processing Status & Results",
                lines=15,
                max_lines=20
            )
    
    # Examples
    gr.Examples(
        examples=[
            ["https://drive.google.com/file/d/1VWEIBdBDVndR-bCiawh29y66lhMkDw5D/view?usp=sharing"]
        ],
        inputs=video_url
    )
    
    # Handle submission
    submit_btn.click(
        fn=process_video,
        inputs=video_url,
        outputs=output
    )

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