import gradio as gr

from inference import generate_video

DESCRIPTION = """
# 🎬 Wan2.1 Text-to-Video

Generate cinematic AI videos with Wan2.1-T2V-1.3B running on Hugging Face ZeroGPU.

### Features

- ⚡ ZeroGPU
- 🎥 480P Output
- ⏱️ Adjustable Video Length
- 🎲 Random Seed
- 📥 MP4 Download
"""

CSS = """
footer {
    visibility: hidden;
}

.gradio-container {
    max-width: 1200px !important;
    margin: auto;
}

#generate-btn {
    height: 56px;
    font-size: 18px;
    font-weight: 600;
}

video {
    border-radius: 12px;
}
"""

EXAMPLES = [
    ["A cinematic drone shot over snowy mountains during sunrise"],
    ["A futuristic cyberpunk street at night with neon lights and rain"],
    ["A golden retriever surfing giant ocean waves"],
    ["An astronaut walking through an ancient Egyptian temple"],
]


def update_length(seconds):
    fps = 16
    frames = int(seconds * fps)

    # Wan requires 4k + 1
    frames = ((frames - 1) // 4) * 4 + 1
    frames = max(33, min(frames, 161))

    return f"Estimated frames: **{frames}**"


with gr.Blocks(title="Wan2.1 Text-to-Video") as demo:

    gr.Markdown(DESCRIPTION)

    with gr.Row():

        with gr.Column(scale=2):

            prompt = gr.Textbox(
                label="Prompt",
                lines=5,
                placeholder="Describe the video you want to generate..."
            )

            negative_prompt = gr.Textbox(
                label="Negative Prompt",
                lines=3,
                value="low quality, blurry, watermark, logo, text, distorted"
            )

            gr.Examples(
                examples=EXAMPLES,
                inputs=prompt,
            )

            length = gr.Slider(
                minimum=2,
                maximum=10,
                value=5,
                step=1,
                label="Video Length (seconds)"
            )

            frame_info = gr.Markdown("Estimated frames: **81**")

            length.change(
                fn=update_length,
                inputs=length,
                outputs=frame_info,
            )

            with gr.Accordion("Advanced Settings", open=False):

                steps = gr.Slider(
                    minimum=20,
                    maximum=60,
                    value=30,
                    step=1,
                    label="Inference Steps",
                )

                guidance = gr.Slider(
                    minimum=1,
                    maximum=10,
                    value=5,
                    step=0.5,
                    label="Guidance Scale",
                )

                seed = gr.Number(
                    value=-1,
                    precision=0,
                    label="Seed (-1 = Random)",
                )

            generate_btn = gr.Button(
                "🎬 Generate Video",
                variant="primary",
                elem_id="generate-btn",
            )

        with gr.Column(scale=1):

            output_video = gr.Video(
                label="Generated Video",
                autoplay=True,
            )

            status = gr.Textbox(
                label="Status",
                interactive=False,
            )

    generate_btn.click(
        fn=generate_video,
        inputs=[
            prompt,
            negative_prompt,
            steps,
            guidance,
            length,
            seed,
        ],
        outputs=[
            output_video,
            status,
        ],
        concurrency_limit=1,
    )


if __name__ == "__main__":
    demo.queue(
        default_concurrency_limit=1,
        max_size=20,
    ).launch(
        theme=gr.themes.Soft(),
        css=CSS,
    )