import gradio as gr
import os
import tempfile
import shutil
from typing import Optional, Union
from huggingface_hub import InferenceClient
from pathlib import Path

# Initialize Hugging Face Inference Client with fal-ai provider
client = InferenceClient(
    provider="fal-ai",
    api_key=os.environ.get("HF_TOKEN"),
    bill_to="huggingface",
)

def cleanup_temp_files():
    """Clean up old temporary video files to prevent storage overflow."""
    try:
        temp_dir = tempfile.gettempdir()
        # Clean up old .mp4 files in temp directory
        for file_path in Path(temp_dir).glob("*.mp4"):
            try:
                # Remove files older than 5 minutes
                import time
                if file_path.stat().st_mtime < (time.time() - 300):
                    file_path.unlink(missing_ok=True)
            except Exception:
                pass
    except Exception as e:
        print(f"Cleanup error: {e}")

def generate_video(
    prompt: str,
    duration: int = 8,
    size: str = "1280x720",
    api_key: Optional[str] = None
) -> Optional[str]:
    """Generate video using Sora-2 through Hugging Face Inference API with fal-ai provider."""
    cleanup_temp_files()
    try:
        if api_key:
            temp_client = InferenceClient(
                provider="fal-ai",
                api_key=api_key,
                bill_to="huggingface",
            )
        else:
            temp_client = client
            if not os.environ.get("HF_TOKEN") and not api_key:
                return None
        
        video_bytes = temp_client.text_to_video(
            prompt,
            model="akhaliq/sora-2",
        )
        
        temp_file = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
        try:
            temp_file.write(video_bytes)
            temp_file.flush()
            video_path = temp_file.name
        finally:
            temp_file.close()
        
        return video_path
    except Exception as e:
        return None

def generate_video_from_image(
    image: Union[str, bytes],
    prompt: str,
    api_key: Optional[str] = None
) -> Optional[str]:
    """Generate a video from a single input image + prompt using Sora-2 image-to-video."""
    cleanup_temp_files()
    if not prompt or prompt.strip() == "":
        return None
    try:
        if api_key:
            temp_client = InferenceClient(
                provider="fal-ai",
                api_key=api_key,
                bill_to="huggingface",
            )
        else:
            temp_client = client
            if not os.environ.get("HF_TOKEN") and not api_key:
                return None

        if isinstance(image, str):
            with open(image, "rb") as f:
                input_image = f.read()
        elif isinstance(image, (bytes, bytearray)):
            input_image = image
        else:
            return None

        video_bytes = temp_client.image_to_video(
            input_image,
            prompt=prompt,
            model="akhaliq/sora-2-image-to-video",
        )

        temp_file = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
        try:
            temp_file.write(video_bytes)
            temp_file.flush()
            video_path = temp_file.name
        finally:
            temp_file.close()

        return video_path
    except Exception as e:
        return None

def generate_with_auth(
    prompt: str, 
    profile: gr.OAuthProfile | None
) -> Optional[str]:
    """Wrapper function that checks if user is logged in before generating video."""
    if profile is None:
        raise gr.Error("Click Sign in with Hugging Face button to use this app for free")
    
    if not prompt or prompt.strip() == "":
        return None
    
    return generate_video(
        prompt, 
        duration=8, 
        size="1280x720", 
        api_key=None
    )

def generate_with_auth_image(
    prompt: str,
    image_path: Optional[str],
    profile: gr.OAuthProfile | None
) -> Optional[str]:
    """Checks login status then calls image->video generator."""
    if profile is None:
        raise gr.Error("Click Sign in with Hugging Face button to use this app for free")
    if not image_path:
        return None
    return generate_video_from_image(image=image_path, prompt=prompt, api_key=None)

def create_ui():
    css = '''
    .logo-dark{display: none}
    .dark .logo-dark{display: block !important}
    .dark .logo-light{display: none}
    #sub_title{margin-top: -20px !important}
    .mobile-notice {
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        color: white;
        padding: 15px;
        border-radius: 10px;
        margin: 20px auto;
        max-width: 600px;
        text-align: center;
    }
    .mobile-notice a {
        color: #fff;
        text-decoration: underline;
        font-weight: bold;
    }
    '''
    
    with gr.Blocks(title="Sora-2", theme=gr.themes.Soft(), css=css) as demo:
        gr.HTML("""
            <div style="text-align: center; max-width: 800px; margin: 0 auto;">
                <h1 style="font-size: 2.5em; margin-bottom: 0.5em;">
                    🎬 Sora-2
                </h1>
                <p style="font-size: 1.1em; color: #666; margin-bottom: 20px;">Generate stunning videos using OpenAI's Sora-2 model</p>
                <div class="mobile-notice">
                    📱 On mobile? Use the optimized version: <a href="https://akhaliq-sora-2.hf.space" target="_blank">akhaliq-sora-2.hf.space</a>
                </div>
                <p style='color: orange;'>⚠️ You must Sign in with Hugging Face using the button to use this app.</p>
                <p style="font-size: 0.9em; color: #999; margin-top: 15px;">
                    Built with <a href="https://huggingface.co/spaces/akhaliq/anycoder" target="_blank" style="color: #667eea;">anycoder</a>
                </p>
            </div>
        """)
        
        # Add login button - required for OAuth
        gr.LoginButton()
        
        # Text -> Video
        with gr.Row():
            with gr.Column(scale=1):
                prompt_input = gr.Textbox(
                    label="Enter your prompt",
                    placeholder="Describe the video you want to create...",
                    lines=4,
                    elem_id="prompt-text-input"
                )
                generate_btn = gr.Button("🎥 Generate Video", variant="primary", size="lg")
            with gr.Column(scale=1):
                video_output = gr.Video(
                    label="Generated Video", 
                    height=400, 
                    interactive=False, 
                    show_download_button=True,
                    elem_id="text-to-video"
                )
        
        generate_btn.click(
            fn=generate_with_auth,
            inputs=[prompt_input],
            outputs=[video_output],
        )

        # Image -> Video UI
        gr.HTML("""
            <div style="text-align: center; margin: 40px 0 10px;">
                <h3 style="margin-bottom: 8px;">🖼️ ➜ 🎬 Image → Video (beta)</h3>
                <p style="color:#666; margin:0;">Turn a single image into a short video with a guiding prompt.</p>
            </div>
        """)
        with gr.Row():
            with gr.Column(scale=1):
                img_prompt_input = gr.Textbox(
                    label="Describe how the scene should evolve",
                    placeholder="e.g., The cat starts to dance and spins playfully",
                    lines=3,
                    elem_id="img-prompt-text-input"
                )
                image_input = gr.Image(label="Upload an image", type="filepath")
                generate_img_btn = gr.Button("🎥 Generate from Image", variant="primary")
            with gr.Column(scale=1):
                video_output_img = gr.Video(
                    label="Generated Video (from Image)", 
                    height=400, 
                    interactive=False, 
                    show_download_button=True,
                    elem_id="image-to-video"
                )

        generate_img_btn.click(
            fn=generate_with_auth_image,
            inputs=[img_prompt_input, image_input],
            outputs=[video_output_img],
        )
        
        # Example usage guidance
        gr.Examples(
            examples=[
                ["A majestic golden eagle soaring through a vibrant sunset sky"],
            ],
            inputs=prompt_input,
            outputs=video_output,
            fn=generate_video,
            cache_examples=False,
            api_name=False,
            show_api=False,
        )
    
    return demo

if __name__ == "__main__":
    try:
        cleanup_temp_files()
        if os.path.exists("gradio_cached_examples"):
            shutil.rmtree("gradio_cached_examples", ignore_errors=True)
    except Exception as e:
        print(f"Initial cleanup error: {e}")
    
    app = create_ui()
    # Configure queue with optimized settings for OAuth-enabled app
    app.queue(
        status_update_rate="auto",
        api_open=False,  # Disable public API access for security
        default_concurrency_limit=None  # Allow multiple concurrent requests
    )
    app.launch(
        show_api=False,
        enable_monitoring=False,
        quiet=True,
        ssr_mode=True
    )