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

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

# 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 initialize_sdk_config():
    """Initialize SDK configuration as per docs"""
    try:
        configuration = byteplussdkcore.Configuration()
        configuration.client_side_validation = True
        configuration.schema = "http"
        configuration.debug = False
        configuration.logger_file = "sdk.log"
        byteplussdkcore.Configuration.set_default(configuration)
        return True
    except Exception as e:
        logger.error(f"SDK config error: {e}")
        return False

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...")
        # Initialize SDK config
        initialize_sdk_config()
        
        # Set environment variable as per original code
        os.environ["ARK_API_KEY"] = api_key
        
        # Initialize client
        client = Ark(
            base_url="https://ark.ap-southeast.bytepluses.com/api/v3",
            api_key=os.environ.get("ARK_API_KEY"),
        )
        
        progress(0.1, desc="Creating video generation request...")
        logger.info("🚀 Creating video generation request...")
        
        # Create task as per original code
        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 as per original code
        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!")
                # Try to 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:
                    # If no URL, show the full result for debugging
                    result_str = str(get_result)
                    yield f"✅ Task completed. Response: {result_str[:500]}...", 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

# Custom CSS for better UI
custom_css = """
.gradio-container {
    max-width: 1200px !important;
    margin: auto !important;
}
.video-container {
    border-radius: 8px;
    overflow: hidden;
}
.status-box {
    background: #f5f5f5;
    border-radius: 8px;
    padding: 10px;
    margin: 10px 0;
}
"""

# Create Gradio interface with latest features
with gr.Blocks(
    title="BytePlus Video Generator",
    theme=gr.themes.Soft(
        primary_hue="blue",
        secondary_hue="purple",
    ),
    css=custom_css
) as demo:
    
    gr.Markdown("""
    # 🎥 BytePlus Video Generator
    Generate stunning videos from images and text prompts using BytePlus AI
    
    ### 📊 System Status
    """)
    
    # SDK Status with modern indicator
    with gr.Row():
        if SDK_AVAILABLE:
            gr.Markdown("✅ **SDK Status:** Connected to BytePlus SDK")
        else:
            gr.Markdown("❌ **SDK Status:** SDK not available")
    
    gr.Markdown("---")
    
    with gr.Row():
        with gr.Column(scale=1, variant="panel"):
            # API Key input with better UX
            api_key = gr.Textbox(
                label="🔑 BytePlus API Key",
                placeholder="Enter your BytePlus API key here",
                type="password",
                value="key",
                info="Your API key will be set as ARK_API_KEY environment variable",
                container=True,
                scale=1
            )
            
            # Model ID with examples
            model_id = gr.Dropdown(
                label="🤖 Model Selection",
                choices=[
                    "seedance-1-5-pro-251215",
                    "seedance-1-5-pro-251215-v2",
                    "byteplus-video-v1"
                ],
                value="seedance-1-5-pro-251215",
                info="Select the model for video generation",
                allow_custom_value=True
            )
            
            # Image URL input with preview
            with gr.Group():
                image_url_input = gr.Textbox(
                    label="🔗 Image URL",
                    placeholder="Enter public image URL",
                    value="https://ark-doc.tos-ap-southeast-1.bytepluses.com/seepro_i2v%20.png",
                    info="Image must be publicly accessible"
                )
                
                image_preview = gr.Image(
                    label="Image Preview",
                    type="pil",
                    height=200,
                    interactive=False
                )
            
            # Text prompt with character count
            prompt_input = gr.Textbox(
                label="📝 Text Prompt",
                placeholder="Describe your video...",
                value="At breakneck speed, drones thread through intricate obstacles or stunning natural wonders, delivering an immersive, heart-pounding flying experience. --duration 5 --camerafixed false",
                lines=4,
                max_lines=6,
                show_copy_button=True
            )
            
            # Generate button with loading state
            generate_btn = gr.Button(
                "🚀 Generate Video",
                variant="primary",
                size="lg"
            )
        
        with gr.Column(scale=1, variant="panel"):
            # Status output with better styling
            status_output = gr.Textbox(
                label="📊 Generation Status",
                lines=5,
                show_copy_button=True,
                container=True
            )
            
            # Video output with player
            with gr.Group(elem_classes="video-container"):
                video_output = gr.Video(
                    label="🎬 Generated Video",
                    interactive=False,
                    show_download_button=True,
                    show_share_button=True
                )
            
            # Video URL with copy button
            video_url_output = gr.Textbox(
                label="🔗 Video URL",
                interactive=False,
                show_copy_button=True
            )
    
    # Example prompts section with gallery
    gr.Markdown("---")
    gr.Markdown("## 📋 Example Prompts")
    
    with gr.Row():
        with gr.Column():
            example1 = gr.Button("🌄 Nature Drone", size="sm")
            gr.Markdown("Aerial shot over mountains at sunrise")
        with gr.Column():
            example2 = gr.Button("🏙️ City Racing", size="sm")
            gr.Markdown("Fast drone racing through neon city")
        with gr.Column():
            example3 = gr.Button("🌊 Ocean Waves", size="sm")
            gr.Markdown("Drone following surfers on waves")
    
    # Function to update image preview
    def update_preview(url):
        try:
            if url and url.startswith(('http://', 'https://')):
                from PIL import Image
                import requests
                from io import BytesIO
                
                response = requests.get(url, timeout=5)
                img = Image.open(BytesIO(response.content))
                return img
            return None
        except:
            return None
    
    # Event handlers
    def set_nature():
        return "Aerial drone shot flying over majestic mountains at sunrise, cinematic lighting, smooth motion --duration 5 --camerafixed false"
    
    def set_city():
        return "Fast-paced drone racing through futuristic city streets with neon lights, dynamic angles, high speed --duration 5 --camerafixed false"
    
    def set_ocean():
        return "Drone following surfers riding massive waves, slow motion, dramatic ocean views, golden hour --duration 5 --camerafixed false"
    
    # Connect events
    image_url_input.change(
        fn=update_preview,
        inputs=image_url_input,
        outputs=image_preview
    )
    
    example1.click(fn=set_nature, outputs=prompt_input)
    example2.click(fn=set_city, outputs=prompt_input)
    example3.click(fn=set_ocean, outputs=prompt_input)
    
    # Generate with progress bar
    generate_event = generate_btn.click(
        fn=generate_video,
        inputs=[api_key, prompt_input, image_url_input, model_id],
        outputs=[status_output, video_output],
        show_progress='full'
    )
    
    # Update video URL when video changes
    generate_event.then(
        fn=lambda url: url if url else "No URL available",
        inputs=[video_output],
        outputs=[video_url_output]
    )
    
    # Clear button with confirmation
    clear_btn = gr.Button("🗑️ Clear All", variant="secondary", size="sm")
    clear_btn.click(
        fn=lambda: (
            "https://ark-doc.tos-ap-southeast-1.bytepluses.com/seepro_i2v%20.png",
            "Enter your prompt here...",
            "",
            None,
            ""
        ),
        outputs=[image_url_input, prompt_input, status_output, video_output, video_url_output]
    )
    
    # Add footer
    gr.Markdown("---")
    gr.Markdown("""
    <div style='text-align: center'>
        <p>Powered by BytePlus SDK | Updated with Gradio 5.23.3</p>
        <p style='font-size: 0.8em; color: gray;'>Images must be publicly accessible. Generation takes 30-60 seconds.</p>
    </div>
    """)

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