import gradio as gr
import os
import time
import json
from PIL import Image
import logging
import sys

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

# Print Gradio version for debugging
logger.info(f"Gradio version: {gr.__version__}")

# Import BytePlus SDK
try:
    import byteplussdkcore
    from byteplussdkarkruntime import Ark
    SDK_AVAILABLE = True
    logger.info("✅ BytePlus SDK loaded successfully")
except ImportError as e:
    SDK_AVAILABLE = False
    logger.warning(f"⚠️ BytePlus SDK not available: {e}")

def generate_video(api_key, prompt_text, image_url, model_id, progress=gr.Progress()):
    """Generate video using the BytePlus SDK"""
    
    if not api_key or api_key == "key":
        yield "⚠️ Please enter your actual BytePlus API key (the 'key' is just a placeholder)", None
        return
    
    if not SDK_AVAILABLE:
        yield "❌ BytePlus SDK not available. Please check installation of byteplus-python-sdk-v2", None
        return
    
    try:
        progress(0, desc="Initializing SDK...")
        
        # Set environment variable
        os.environ["Key"] = api_key
        
        # Initialize client
        client = Ark(
            base_url="https://ark.ap-southeast.bytepluses.com/api/v3",
            api_key=os.environ.get("Key"),
        )
        
        progress(0.1, desc="Creating video generation request...")
        logger.info("🚀 Creating video generation request...")
        
        # Create task
        create_result = client.content_generation.tasks.create(
            model=model_id,
            content=[
                {
                    "type": "text",
                    "text": prompt_text
                },
                {
                    "type": "image_url",
                    "image_url": {
                        "url": image_url
                    }
                }
            ]
        )
        
        task_id = create_result.id
        logger.info(f"📋 Task ID: {task_id}")
        
        yield f"⏳ Task created with ID: {task_id}. Waiting for completion...", None
        
        # Polling
        max_attempts = 60
        attempts = 0
        
        while attempts < max_attempts:
            progress(0.1 + (attempts / max_attempts) * 0.8, 
                    desc=f"Polling status: {attempts + 1}/{max_attempts}")
            
            get_result = client.content_generation.tasks.get(task_id=task_id)
            status = get_result.status
            
            if status == "succeeded":
                progress(1.0, desc="Complete!")
                logger.info("✅ Task succeeded!")
                
                # Extract video URL
                video_url = None
                if hasattr(get_result, 'output') and get_result.output:
                    if isinstance(get_result.output, list) and len(get_result.output) > 0:
                        video_url = get_result.output[0].get('video_url')
                    elif hasattr(get_result.output, 'video_url'):
                        video_url = get_result.output.video_url
                    elif isinstance(get_result.output, dict):
                        video_url = get_result.output.get('video_url')
                
                if video_url:
                    yield f"✅ Video generated successfully!", video_url
                else:
                    yield f"✅ Task completed but no video URL in response", None
                break
                
            elif status == "failed":
                error_msg = get_result.error if hasattr(get_result, 'error') else "Unknown error"
                yield f"❌ Task failed: {error_msg}", None
                break
            else:
                yield f"⏳ Current status: {status}, waiting... ({attempts + 1}/{max_attempts})", None
                time.sleep(1)
                attempts += 1
        
        if attempts >= max_attempts:
            yield "⏰ Timeout: Task took too long to complete. Please try again.", None
            
    except Exception as e:
        logger.error(f"Error: {e}")
        yield f"❌ Error: {str(e)}", None

# Simple, clean UI
with gr.Blocks(title="BytePlus Video Generator", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🎥 BytePlus Video Generator
    Generate videos from images and text prompts
    """)
    
    with gr.Row():
        with gr.Column():
            api_key = gr.Textbox(
                label="🔑 API Key",
                type="password",
                placeholder="Enter your BytePlus API key",
                value="key"
            )
            
            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=3,
                value="At breakneck speed, drones thread through intricate obstacles --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")
    
    generate_btn.click(
        fn=generate_video,
        inputs=[api_key, prompt, image_url, model_id],
        outputs=[status, video]
    ).then(
        fn=lambda url: url if url else "",
        inputs=[video],
        outputs=[video_url]
    )

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