import gradio as gr
import os
from huggingface_hub import InferenceClient
import tempfile
from pathlib import Path

# Initialize the client
client = InferenceClient(
    provider="fal-ai",
    api_key=os.environ.get("HF_TOKEN"),
    bill_to="huggingface",
)

def text_to_video(prompt, aspect_ratio="16:9", profile: gr.OAuthProfile | None = None):
    """Generate video from text prompt"""
    try:
        if profile is None:
            return None, "❌ Click Sign in with Hugging Face button to use this app for free"

        if not prompt or prompt.strip() == "":
            return None, "Please enter a text prompt"

        # Generate video from text with aspect ratio parameter
        video = client.text_to_video(
            prompt,
            model="akhaliq/veo3.1-fast",
            aspect_ratio=aspect_ratio,
        )

        # Save the video
        with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as tmp_file:
            tmp_file.write(video)
            video_path = tmp_file.name

        return video_path, f"✅ Video generated successfully from prompt: '{prompt[:50]}...'"

    except Exception as e:
        return None, f"❌ Error generating video: {str(e)}"


def image_to_video(image, prompt, aspect_ratio="16:9", profile: gr.OAuthProfile | None = None):
    """Generate video from image and prompt"""
    try:
        if profile is None:
            return None, "❌ Click Sign in with Hugging Face button to use this app for free"

        if image is None:
            return None, "Please upload an image"

        if not prompt or prompt.strip() == "":
            return None, "Please enter a prompt describing the motion"

        import io
        from PIL import Image as PILImage

        # Convert image to bytes
        buffer = io.BytesIO()
        if isinstance(image, PILImage.Image):
            image.save(buffer, format='PNG')
        else:
            PILImage.fromarray(image).save(buffer, format='PNG')
        input_image = buffer.getvalue()

        # Generate video from image with aspect ratio parameter
        video = client.image_to_video(
            input_image,
            prompt=prompt,
            model="akhaliq/veo3.1-fast-image-to-video",
            aspect_ratio=aspect_ratio,
        )

        # Save the video
        with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as tmp_file:
            tmp_file.write(video)
            video_path = tmp_file.name

        return video_path, f"✅ Video generated successfully with motion: '{prompt[:50]}...'"

    except Exception as e:
        return None, f"❌ Error generating video: {str(e)}"


# Custom CSS
custom_css = """
.container {
    max-width: 1200px;
    margin: auto;
}
.status-box {
    padding: 10px;
    border-radius: 5px;
    margin-top: 10px;
}
"""

# Interface
with gr.Blocks(css=custom_css, theme=gr.themes.Soft(), title="AI Video Generator") as demo:
    gr.Markdown("# 🎬 AI Video Generator\n### Generate videos from text or animate images with AI")

    gr.LoginButton()

    with gr.Tabs() as tabs:
        # --- TEXT TO VIDEO TAB ---
        with gr.Tab("📝 Text to Video", id=0):
            text_prompt = gr.Textbox(
                label="Text Prompt",
                placeholder="Describe the video you want to create...",
                lines=4
            )
            text_aspect = gr.Radio(
                choices=["16:9", "9:16"],
                value="16:9",
                label="Aspect Ratio"
            )

            with gr.Row():
                text_generate_btn = gr.Button("🎬 Generate Video", variant="primary")
                text_clear_btn = gr.ClearButton(value="🗑️ Clear")

            text_status = gr.Textbox(label="Status", interactive=False)
            text_video_output = gr.Video(label="Generated Video", autoplay=True, show_download_button=True)

            text_generate_btn.click(
                fn=text_to_video,
                inputs=[text_prompt, text_aspect],
                outputs=[text_video_output, text_status],
                show_progress="full"
            )

            text_clear_btn.click(
                fn=lambda: ("", None, ""),
                outputs=[text_prompt, text_video_output, text_status]
            )

        # --- IMAGE TO VIDEO TAB ---
        with gr.Tab("🖼️ Image to Video", id=1):
            image_input = gr.Image(label="Upload Image", type="pil", height=300)
            image_prompt = gr.Textbox(label="Motion Prompt", placeholder="Describe how the image should move...")
            image_aspect = gr.Radio(
                choices=["16:9", "9:16"],
                value="16:9",
                label="Aspect Ratio"
            )

            with gr.Row():
                image_generate_btn = gr.Button("🎬 Animate Image", variant="primary")
                image_clear_btn = gr.ClearButton(value="🗑️ Clear")

            image_status = gr.Textbox(label="Status", interactive=False)
            image_video_output = gr.Video(label="Generated Video", autoplay=True, show_download_button=True)

            image_generate_btn.click(
                fn=image_to_video,
                inputs=[image_input, image_prompt, image_aspect],
                outputs=[image_video_output, image_status],
                show_progress="full"
            )

            image_clear_btn.click(
                fn=lambda: (None, "", None, ""),
                outputs=[image_input, image_prompt, image_video_output, image_status]
            )

# Launch
if __name__ == "__main__":
    demo.launch(share=False, show_api=False, show_error=True, quiet=True)