import os
import torch
import gradio as gr
from diffusers import AutoencoderKLWan, WanPipeline
from diffusers.utils import export_to_video

# Device aur dtype configuration
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.bfloat16 if device == "cuda" else torch.float32

print(f"Initializing LYRA AI Generator on device: {device}")

_pipeline = None

def get_wan_pipeline():
    global _pipeline
    if _pipeline is None:
        model_id = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
        print("Loading Wan2.1-T2V-1.3B model... Please wait.")
        vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float32)
        _pipeline = WanPipeline.from_pretrained(model_id, vae=vae, torch_dtype=dtype)
        _pipeline.to(device)
    return _pipeline

def generate_video(prompt, negative_prompt, guidance_scale):
    if not prompt.strip():
        raise gr.Error("Kripya ek valid prompt enter karein!")
    
    try:
        pipe = get_wan_pipeline()
        num_frames = 150  # ~10 seconds at 15 fps
        
        print(f"Generating video for prompt: '{prompt}'...")
        output_frames = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            height=480,
            width=832,
            num_frames=num_frames,
            guidance_scale=guidance_scale
        ).frames[0]
        
        video_path = "lyra_output.mp4"
        export_to_video(output_frames, video_path, fps=15)
        print("Video generated successfully!")
        return video_path

    except Exception as e:
        print(f"Error during generation: {e}")
        raise gr.Error(f"Generation error: {str(e)}. Ensure you are running on a GPU-enabled space.")

# Modern Dark Glassmorphism CSS Styling
custom_css = """
body {
    background: linear-gradient(135deg, #090a0f 0%, #131722 50%, #1f1a3a 100%) !important;
    color: #f3f4f6 !important;
    font-family: 'Inter', sans-serif;
}
.gradio-container {
    max-width: 1100px !important;
    margin: auto;
    border-radius: 20px;
    border: 1px solid rgba(255, 255, 255, 0.1);
    background: rgba(15, 23, 42, 0.75) !important;
    backdrop-filter: blur(16px);
    box-shadow: 0 20px 50px rgba(0, 0, 0, 0.5);
}
.header-title {
    text-align: center;
    background: linear-gradient(90deg, #c084fc, #f472b6, #38bdf8);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    font-weight: 800;
    font-size: 2.8rem;
    margin-bottom: 0px;
}
.header-subtitle {
    text-align: center;
    color: #94a3b8;
    font-size: 1.1rem;
    margin-top: 5px;
    margin-bottom: 25px;
}
.creator-badge {
    text-align: center;
    color: #38bdf8;
    font-size: 0.95rem;
    font-weight: 600;
    letter-spacing: 1px;
    text-transform: uppercase;
    margin-bottom: 20px;
}
.generate-btn {
    background: linear-gradient(135deg, #6366f1 0%, #a855f7 50%, #ec4899 100%) !important;
    border: none !important;
    color: white !important;
    font-weight: 700 !important;
    font-size: 1.1rem !important;
    border-radius: 12px !important;
    transition: all 0.3s ease !important;
    box-shadow: 0 4px 15px rgba(168, 85, 247, 0.4);
}
.generate-btn:hover {
    transform: translateY(-2px);
    box-shadow: 0 6px 20px rgba(236, 72, 153, 0.6);
}
"""

with gr.Blocks(css=custom_css, theme=gr.themes.Soft(primary_hue="purple", secondary_hue="pink")) as demo:
    gr.HTML("""
        <div class="header-title">✨ LYRA AI VIDEO GENERATOR ✨</div>
        <div class="header-subtitle">State-of-the-art Text-to-Video AI Studio powered by Wan2.1</div>
        <div class="creator-badge">🚀 Created by Muhammad Taqi</div>
    """)
    
    with gr.Row(equal_height=True):
        with gr.Column(scale=1):
            prompt_input = gr.Textbox(
                label="🎬 Video Prompt",
                placeholder="A cinematic drone shot of a futuristic neon cyberpunk city...",
                lines=4,
                value="A cinematic drone shot of a futuristic cyberpunk city with neon lights reflecting on wet streets, 4k."
            )
            
            with gr.Accordion("Advanced Options", open=False):
                negative_prompt_input = gr.Textbox(
                    label="❌ Negative Prompt",
                    value="Bright tones, overexposed, static, blurred details, low quality",
                    lines=2
                )
                guidance_slider = gr.Slider(
                    minimum=1.0, maximum=10.0, value=5.0, step=0.5,
                    label="Guidance Scale"
                )
            
            submit_btn = gr.Button("🔮 Generate Masterpiece (10-20s)", elem_classes="generate-btn")
            
        with gr.Column(scale=1):
            video_output = gr.Video(
                label="🎥 Generated Video Result (Downloadable)",
                format="mp4",
                autoplay=True,
                show_share_button=True
            )
            gr.Markdown("*Tip: Click the download button on the video player to save your file!*")

    submit_btn.click(
        fn=generate_video,
        inputs=[prompt_input, negative_prompt_input, guidance_slider],
        outputs=video_output
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)