import gradio as gr
import requests
import os
import time
import json
import logging

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Get API key from environment (your "Key" secret)
API_KEY = os.environ.get("Key", "")
if not API_KEY:
    logger.error("❌ No API key found in environment variable 'Key'")
else:
    logger.info(f"✅ API key loaded (length: {len(API_KEY)})")

def generate_video(prompt_text, image_url, model_id, progress=gr.Progress()):
    """Generate video using direct API calls (curl approach)"""
    
    if not API_KEY:
        yield "❌ API key not configured. Please set the 'Key' secret.", None
        return
    
    try:
        progress(0, desc="Starting request...")
        
        # Use the path from curl command
        url = "https://ark.ap-southeast.bytepluses.com/api/v3/contents/generations/tasks"
        
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {API_KEY}"  # Same as curl: -H "Authorization: Bearer $ARK_API_KEY"
        }
        
        # Prepare the request body exactly like curl
        data = {
            "model": model_id,
            "content": [
                {
                    "type": "text",
                    "text": prompt_text
                },
                {
                    "type": "image_url",
                    "image_url": {
                        "url": image_url
                    }
                }
            ]
        }
        
        logger.info(f"📤 Sending request to: {url}")
        logger.info(f"📝 Model: {model_id}")
        
        # Make the POST request (like curl -X POST)
        response = requests.post(url, headers=headers, json=data)
        
        logger.info(f"📥 Response status: {response.status_code}")
        
        if response.status_code != 200:
            error_msg = f"Error {response.status_code}: {response.text}"
            logger.error(f"❌ {error_msg}")
            yield error_msg, None
            return
        
        result = response.json()
        task_id = result.get('id')
        
        if not task_id:
            yield "❌ No task ID in response", None
            return
        
        logger.info(f"✅ Task created: {task_id}")
        yield f"Task created: {task_id}", None
        
        # Poll for results (like the while loop in original code)
        progress(0.2, desc="Waiting for generation...")
        
        # Status URL (adjust if different)
        status_url = f"https://ark.ap-southeast.bytepluses.com/api/v3/contents/generations/tasks/{task_id}"
        
        max_attempts = 60
        attempts = 0
        
        while attempts < max_attempts:
            progress(0.2 + (attempts / max_attempts) * 0.7, 
                    desc=f"Polling... ({attempts + 1}/{max_attempts})")
            
            status_response = requests.get(status_url, headers=headers)
            
            if status_response.status_code != 200:
                yield f"❌ Error checking status: {status_response.text}", None
                return
            
            status_result = status_response.json()
            status = status_result.get('status')
            
            logger.info(f"Status: {status}")
            
            if status == "succeeded":
                progress(1.0, desc="Complete!")
                
                # Extract video URL (adjust based on actual response structure)
                video_url = None
                if 'output' in status_result and status_result['output']:
                    if isinstance(status_result['output'], list) and len(status_result['output']) > 0:
                        video_url = status_result['output'][0].get('video_url')
                    elif isinstance(status_result['output'], dict):
                        video_url = status_result['output'].get('video_url')
                
                if video_url:
                    yield "✅ Video generated successfully!", video_url
                else:
                    yield "✅ Task completed (no video URL in response)", None
                return
                
            elif status == "failed":
                error = status_result.get('error', 'Unknown error')
                yield f"❌ Task failed: {error}", None
                return
            else:
                yield f"⏳ Status: {status}...", None
                time.sleep(1)
                attempts += 1
        
        yield "⏰ Timeout: Task took too long", None
        
    except Exception as e:
        logger.error(f"❌ Exception: {str(e)}")
        yield f"❌ Error: {str(e)}", None

# Simple Gradio interface
with gr.Blocks(title="BytePlus Video Generator", theme=gr.themes.Soft()) as demo:
    
    gr.Markdown("""
    # 🎥 BytePlus Video Generator
    
    Using direct API calls (curl approach)
    """)
    
    # Show API key status
    if API_KEY:
        gr.Markdown("✅ **API Key:** Configured")
    else:
        gr.Markdown("❌ **API Key:** Not configured - please add 'Key' secret")
    
    with gr.Row():
        with gr.Column():
            model_id = gr.Textbox(
                label="Model ID",
                value="seedance-1-5-pro-251215"
            )
            
            image_url = gr.Textbox(
                label="Image URL",
                value="https://ark-doc.tos-ap-southeast-1.bytepluses.com/seepro_i2v%20.png"
            )
            
            prompt = gr.Textbox(
                label="Prompt",
                lines=4,
                value="At breakneck speed, drones thread through intricate obstacles or stunning natural wonders, delivering an immersive, heart-pounding flying experience. --duration 5 --camerafixed false"
            )
            
            generate_btn = gr.Button("🚀 Generate", variant="primary")
        
        with gr.Column():
            status = gr.Textbox(label="Status", lines=3)
            video = gr.Video(label="Generated Video")
            video_url = gr.Textbox(label="Video URL")
    
    # Example prompts
    gr.Markdown("---")
    with gr.Row():
        gr.Button("🌄 Nature").click(
            fn=lambda: "Aerial drone shot over mountains at sunrise, cinematic --duration 5 --camerafixed false",
            outputs=prompt
        )
        gr.Button("🏙️ City").click(
            fn=lambda: "Fast drone racing through futuristic city streets --duration 5 --camerafixed false",
            outputs=prompt
        )
        gr.Button("🌊 Ocean").click(
            fn=lambda: "Drone following waves at golden hour --duration 5 --camerafixed false",
            outputs=prompt
        )
    
    # Generate
    generate_btn.click(
        fn=generate_video,
        inputs=[prompt, image_url, model_id],
        outputs=[status, video]
    ).then(
        fn=lambda v: v if v else "",
        inputs=[video],
        outputs=[video_url]
    )

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