import gradio as gr
import numpy as np
import random
import torch
import spaces
import os

from PIL import Image
from diffusers.pipelines.wan.pipeline_wan import WanPipeline
from diffusers.utils.export_utils import export_to_video

import tempfile
from typing import Optional, Tuple, Any
from groq import Groq


# =========================================================
# API CONFIGURATION
# =========================================================
GROQ_API_KEY = os.getenv("GROQ_API_KEY")
MODEL_ID = os.getenv("MODEL_ID")
HF_TOKEN = os.getenv("HF_TOKEN")

# =========================================================
# MODEL LOADING
# =========================================================
dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"

print("Loading pipeline...")
pipe = WanPipeline.from_pretrained(
    MODEL_ID,
    torch_dtype=dtype,
    token=HF_TOKEN
).to(device)

MAX_SEED = np.iinfo(np.int32).max
FIXED_FPS = 24
MIN_FRAMES = 8
MAX_FRAMES = 81

# =========================================================
# DEFAULT PROMPTS
# =========================================================
default_prompt = "A beautiful sunset over the ocean with gentle waves."
default_negative_prompt = "low quality, worst quality, blurry, distorted, deformed, ugly, bad anatomy"

# =========================================================
# PROMPT ENHANCEMENT SYSTEM PROMPT
# =========================================================
ENHANCE_SYSTEM_PROMPT = """You are a professional video prompt engineer. Your task is to enhance user prompts for AI video generation.

Rules:
1. Add vivid visual details (lighting, colors, textures, atmosphere)
2. Include camera movements (pan, zoom, tracking shot, etc.)
3. Describe motion and dynamics clearly
4. Keep the enhanced prompt concise but detailed (max 150 words)
5. Maintain the original intent of the user's prompt
6. Output ONLY the enhanced prompt, nothing else

Example:
User: "A cat playing"
Enhanced: "A fluffy orange tabby cat playfully batting at a dangling yarn ball, soft afternoon sunlight streaming through a window creating warm golden highlights on its fur, smooth tracking shot following the cat's graceful movements, shallow depth of field with bokeh background, cozy living room setting with warm ambient lighting"
"""


# =========================================================
# PROMPT ENHANCEMENT FUNCTION
# =========================================================
def enhance_prompt(prompt: str) -> str:
    """Enhance the user prompt using Groq LLM API."""
    if not GROQ_API_KEY:
        return prompt + " (API key not configured - using original prompt)"
    
    try:
        client = Groq(api_key=GROQ_API_KEY)
        
        enhanced_text = ""
        completion = client.chat.completions.create(
            model="meta-llama/llama-4-scout-17b-16e-instruct",
            messages=[
                {
                    "role": "system",
                    "content": ENHANCE_SYSTEM_PROMPT
                },
                {
                    "role": "user",
                    "content": f"Enhance this video generation prompt: {prompt}"
                }
            ],
            temperature=0.7,
            max_completion_tokens=512,
            top_p=1,
            stream=True,
            stop=None
        )
        
        for chunk in completion:
            if chunk.choices[0].delta.content:
                enhanced_text += chunk.choices[0].delta.content
        
        return enhanced_text.strip() if enhanced_text.strip() else prompt
        
    except Exception as e:
        print(f"Prompt enhancement error: {e}")
        return prompt


# =========================================================
# HELPER FUNCTIONS
# =========================================================
def get_num_frames(duration_seconds: float) -> int:
    """Calculate number of frames based on duration."""
    return max(MIN_FRAMES, min(MAX_FRAMES, int(duration_seconds * FIXED_FPS)))


# =========================================================
# MAIN GENERATION FUNCTION
# =========================================================
@spaces.GPU
def generate_video(
    prompt: str,
    negative_prompt: str = default_negative_prompt,
    enhance_prompt_option: bool = False,
    duration_seconds: float = 3.0,
    guidance_scale: float = 7.5,
    num_inference_steps: int = 20,
    height: int = 480,
    width: int = 832,
    seed: int = 0,
    randomize_seed: bool = True,
    progress=gr.Progress(track_tqdm=True),
) -> Tuple[str, int, str]:
    """Generate video from text prompt."""
    
    if not prompt.strip():
        raise gr.Error("Please enter a prompt.")
    
    # Enhance prompt if option is enabled
    final_prompt = prompt
    if enhance_prompt_option:
        final_prompt = enhance_prompt(prompt)
        print(f"Enhanced Prompt: {final_prompt}")
    
    # Handle seed
    current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
    generator = torch.Generator(device=device).manual_seed(current_seed)
    
    # Calculate frames
    num_frames = get_num_frames(duration_seconds)
    
    # Generate video
    output_frames = pipe(
        prompt=final_prompt,
        negative_prompt=negative_prompt,
        height=height,
        width=width,
        num_frames=num_frames,
        guidance_scale=guidance_scale,
        num_inference_steps=num_inference_steps,
        generator=generator,
    ).frames[0]
    
    # Export to video file
    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmpfile:
        video_path = tmpfile.name
    export_to_video(output_frames, video_path, fps=FIXED_FPS)
    
    # Build info log
    info_log = f"""✅ VIDEO GENERATION COMPLETE!
{'=' * 50}
🎬 Video Info:
   • Duration: {duration_seconds:.1f} seconds
   • Total Frames: {num_frames}
   • FPS: {FIXED_FPS}
   • Resolution: {width} x {height}
{'=' * 50}
⚙️ Generation Settings:
   • Guidance Scale: {guidance_scale}
   • Inference Steps: {num_inference_steps}
   • Seed: {current_seed}
   • Prompt Enhanced: {'Yes' if enhance_prompt_option else 'No'}
{'=' * 50}
💾 Ready to download!"""
    
    return video_path, current_seed, final_prompt, info_log


# ============================================
# 🎨 Comic Classic Theme - Toon Playground
# ============================================

css = """
/* ===== 🎨 Google Fonts Import ===== */
@import url('https://fonts.googleapis.com/css2?family=Bangers&family=Comic+Neue:wght@400;700&display=swap');

/* ===== 🎨 Comic Classic 배경 - 빈티지 페이퍼 + 도트 패턴 ===== */
.gradio-container {
    background-color: #FEF9C3 !important;
    background-image: 
        radial-gradient(#1F2937 1px, transparent 1px) !important;
    background-size: 20px 20px !important;
    min-height: 100vh !important;
    font-family: 'Comic Neue', cursive, sans-serif !important;
}

/* ===== 허깅페이스 상단 요소 숨김 ===== */
.huggingface-space-header,
#space-header,
.space-header,
[class*="space-header"],
.svelte-1ed2p3z,
.space-header-badge,
.header-badge,
[data-testid="space-header"],
.svelte-kqij2n,
.svelte-1ax1toq,
.embed-container > div:first-child {
    display: none !important;
    visibility: hidden !important;
    height: 0 !important;
    width: 0 !important;
    overflow: hidden !important;
    opacity: 0 !important;
    pointer-events: none !important;
}

/* ===== Footer 완전 숨김 ===== */
footer,
.footer,
.gradio-container footer,
.built-with,
[class*="footer"],
.gradio-footer,
.main-footer,
div[class*="footer"],
.show-api,
.built-with-gradio,
a[href*="gradio.app"],
a[href*="huggingface.co/spaces"] {
    display: none !important;
    visibility: hidden !important;
    height: 0 !important;
    padding: 0 !important;
    margin: 0 !important;
}

/* ===== 메인 컨테이너 ===== */
#col-container { 
    max-width: 1000px; 
    margin: 0 auto; 
}

/* ===== 🎨 헤더 타이틀 - 코믹 스타일 ===== */
.header-text h1 {
    font-family: 'Bangers', cursive !important;
    color: #1F2937 !important;
    font-size: 3.5rem !important;
    font-weight: 400 !important;
    text-align: center !important;
    margin-bottom: 0.5rem !important;
    text-shadow: 
        4px 4px 0px #FACC15,
        6px 6px 0px #1F2937 !important;
    letter-spacing: 3px !important;
    -webkit-text-stroke: 2px #1F2937 !important;
}

/* ===== 🎨 서브타이틀 ===== */
.subtitle {
    text-align: center !important;
    font-family: 'Comic Neue', cursive !important;
    font-size: 1.2rem !important;
    color: #1F2937 !important;
    margin-bottom: 1.5rem !important;
    font-weight: 700 !important;
}

/* ===== 🎨 카드/패널 - 만화 프레임 스타일 ===== */
.gr-panel,
.gr-box,
.gr-form,
.block,
.gr-group {
    background: #FFFFFF !important;
    border: 3px solid #1F2937 !important;
    border-radius: 8px !important;
    box-shadow: 6px 6px 0px #1F2937 !important;
    transition: all 0.2s ease !important;
}

.gr-panel:hover,
.block:hover {
    transform: translate(-2px, -2px) !important;
    box-shadow: 8px 8px 0px #1F2937 !important;
}

/* ===== 🎨 입력 필드 (Textbox) ===== */
textarea, 
input[type="text"], 
input[type="number"] {
    background: #FFFFFF !important;
    border: 3px solid #1F2937 !important;
    border-radius: 8px !important;
    color: #1F2937 !important;
    font-family: 'Comic Neue', cursive !important;
    font-size: 1rem !important;
    font-weight: 700 !important;
    transition: all 0.2s ease !important;
}

textarea:focus, 
input[type="text"]:focus, 
input[type="number"]:focus {
    border-color: #3B82F6 !important;
    box-shadow: 4px 4px 0px #3B82F6 !important;
    outline: none !important;
}

textarea::placeholder {
    color: #9CA3AF !important;
    font-weight: 400 !important;
}

/* ===== 🎨 Primary 버튼 - 코믹 블루 ===== */
.gr-button-primary,
button.primary,
.gr-button.primary {
    background: #3B82F6 !important;
    border: 3px solid #1F2937 !important;
    border-radius: 8px !important;
    color: #FFFFFF !important;
    font-family: 'Bangers', cursive !important;
    font-weight: 400 !important;
    font-size: 1.3rem !important;
    letter-spacing: 2px !important;
    padding: 14px 28px !important;
    box-shadow: 5px 5px 0px #1F2937 !important;
    transition: all 0.1s ease !important;
    text-shadow: 1px 1px 0px #1F2937 !important;
}

.gr-button-primary:hover,
button.primary:hover,
.gr-button.primary:hover {
    background: #2563EB !important;
    transform: translate(-2px, -2px) !important;
    box-shadow: 7px 7px 0px #1F2937 !important;
}

.gr-button-primary:active,
button.primary:active,
.gr-button.primary:active {
    transform: translate(3px, 3px) !important;
    box-shadow: 2px 2px 0px #1F2937 !important;
}

/* ===== 🎨 Secondary 버튼 - 코믹 레드 ===== */
.gr-button-secondary,
button.secondary,
.generate-btn {
    background: #EF4444 !important;
    border: 3px solid #1F2937 !important;
    border-radius: 8px !important;
    color: #FFFFFF !important;
    font-family: 'Bangers', cursive !important;
    font-weight: 400 !important;
    font-size: 1.1rem !important;
    letter-spacing: 1px !important;
    box-shadow: 4px 4px 0px #1F2937 !important;
    transition: all 0.1s ease !important;
    text-shadow: 1px 1px 0px #1F2937 !important;
}

.gr-button-secondary:hover,
button.secondary:hover,
.generate-btn:hover {
    background: #DC2626 !important;
    transform: translate(-2px, -2px) !important;
    box-shadow: 6px 6px 0px #1F2937 !important;
}

.gr-button-secondary:active,
button.secondary:active,
.generate-btn:active {
    transform: translate(2px, 2px) !important;
    box-shadow: 2px 2px 0px #1F2937 !important;
}

/* ===== 🎨 로그 출력 영역 ===== */
.info-log textarea {
    background: #1F2937 !important;
    color: #10B981 !important;
    font-family: 'Courier New', monospace !important;
    font-size: 0.9rem !important;
    font-weight: 400 !important;
    border: 3px solid #10B981 !important;
    border-radius: 8px !important;
    box-shadow: 4px 4px 0px #10B981 !important;
}

/* ===== 🎨 비디오 출력 영역 ===== */
.video-output video {
    border: 4px solid #1F2937 !important;
    border-radius: 8px !important;
    box-shadow: 8px 8px 0px #1F2937 !important;
}

/* ===== 🎨 아코디언 - 말풍선 스타일 ===== */
.gr-accordion {
    background: #FACC15 !important;
    border: 3px solid #1F2937 !important;
    border-radius: 8px !important;
    box-shadow: 4px 4px 0px #1F2937 !important;
}

.gr-accordion-header {
    color: #1F2937 !important;
    font-family: 'Comic Neue', cursive !important;
    font-weight: 700 !important;
    font-size: 1.1rem !important;
}

/* ===== 🎨 이미지 출력 영역 ===== */
.gr-image,
.image-container {
    border: 4px solid #1F2937 !important;
    border-radius: 8px !important;
    box-shadow: 8px 8px 0px #1F2937 !important;
    overflow: hidden !important;
    background: #FFFFFF !important;
}

/* ===== 🎨 라벨 스타일 ===== */
label,
.gr-input-label,
.gr-block-label {
    color: #1F2937 !important;
    font-family: 'Comic Neue', cursive !important;
    font-weight: 700 !important;
    font-size: 1rem !important;
}

span.gr-label {
    color: #1F2937 !important;
}

/* ===== 🎨 체크박스 스타일 ===== */
input[type="checkbox"] {
    accent-color: #3B82F6 !important;
    width: 20px !important;
    height: 20px !important;
}

/* ===== 🎨 슬라이더 스타일 ===== */
input[type="range"] {
    accent-color: #3B82F6 !important;
}

/* ===== 🎨 정보 텍스트 ===== */
.gr-info,
.info {
    color: #6B7280 !important;
    font-family: 'Comic Neue', cursive !important;
    font-size: 0.9rem !important;
}

/* ===== 🎨 프로그레스 바 ===== */
.progress-bar,
.gr-progress-bar {
    background: #3B82F6 !important;
    border: 2px solid #1F2937 !important;
    border-radius: 4px !important;
}

/* ===== 🎨 스크롤바 - 코믹 스타일 ===== */
::-webkit-scrollbar {
    width: 12px;
    height: 12px;
}

::-webkit-scrollbar-track {
    background: #FEF9C3;
    border: 2px solid #1F2937;
}

::-webkit-scrollbar-thumb {
    background: #3B82F6;
    border: 2px solid #1F2937;
    border-radius: 0px;
}

::-webkit-scrollbar-thumb:hover {
    background: #EF4444;
}

/* ===== 🎨 선택 하이라이트 ===== */
::selection {
    background: #FACC15;
    color: #1F2937;
}

/* ===== 🎨 링크 스타일 ===== */
a {
    color: #3B82F6 !important;
    text-decoration: none !important;
    font-weight: 700 !important;
}

a:hover {
    color: #EF4444 !important;
}

/* ===== 🎨 Row/Column 간격 ===== */
.gr-row {
    gap: 1.5rem !important;
}

.gr-column {
    gap: 1rem !important;
}

/* ===== 반응형 조정 ===== */
@media (max-width: 768px) {
    .header-text h1 {
        font-size: 2.2rem !important;
        text-shadow: 
            3px 3px 0px #FACC15,
            4px 4px 0px #1F2937 !important;
    }
    
    .gr-button-primary,
    button.primary {
        padding: 12px 20px !important;
        font-size: 1.1rem !important;
    }
    
    .gr-panel,
    .block {
        box-shadow: 4px 4px 0px #1F2937 !important;
    }
}

/* ===== 🎨 다크모드 비활성화 (코믹은 밝아야 함) ===== */
@media (prefers-color-scheme: dark) {
    .gradio-container {
        background-color: #FEF9C3 !important;
    }
}
"""


# =========================================================
# GRADIO UI - Comic Classic Theme
# =========================================================
with gr.Blocks() as demo:
    
    # CSS 삽입
    gr.HTML(f"<style>{css}</style>")
    
    # Header Title
    gr.Markdown(
        """
        # 🎬 UNCENSORED TEXT TO VIDEO 🎥
        """,
        elem_classes="header-text"
    )
    
    gr.Markdown(
        """
        <p class="subtitle">✨ Transform your ideas into stunning AI-generated videos! 🚀</p>
        """,
    )
    
    with gr.Row(equal_height=False):
        # Left column - Input
        with gr.Column(scale=1, min_width=320):
            prompt_input = gr.Textbox(
                label="✏️ Your Prompt",
                value=default_prompt,
                placeholder="Describe the video you want to create...",
                lines=4
            )
            
            enhance_prompt_checkbox = gr.Checkbox(
                label="✨ Enhance Prompt with AI",
                value=False,
                info="Use AI to automatically enhance your prompt for better results"
            )
            
            duration_slider = gr.Slider(
                label="⏱️ Duration (seconds)",
                minimum=1.0,
                maximum=4.0,
                step=0.5,
                value=3.0
            )
            
            generate_btn = gr.Button(
                "🎬 GENERATE VIDEO! 🚀",
                variant="primary",
                size="lg",
                elem_classes="generate-btn"
            )
            
            with gr.Accordion("⚙️ Advanced Options", open=False):
                negative_prompt_input = gr.Textbox(
                    label="Negative Prompt",
                    value=default_negative_prompt,
                    lines=2
                )
                guidance_scale_slider = gr.Slider(
                    label="Guidance Scale",
                    minimum=1.0,
                    maximum=15.0,
                    step=0.5,
                    value=7.5
                )
                num_inference_steps_slider = gr.Slider(
                    label="Inference Steps",
                    minimum=10,
                    maximum=50,
                    step=1,
                    value=20
                )
                height_slider = gr.Slider(
                    label="Height",
                    minimum=256,
                    maximum=720,
                    step=16,
                    value=480
                )
                width_slider = gr.Slider(
                    label="Width",
                    minimum=256,
                    maximum=1280,
                    step=16,
                    value=832
                )
                seed_slider = gr.Slider(
                    label="Seed",
                    minimum=0,
                    maximum=MAX_SEED,
                    step=1,
                    value=0
                )
                randomize_seed_checkbox = gr.Checkbox(
                    label="Randomize Seed",
                    value=True
                )
            
            with gr.Accordion("📜 Generation Info", open=True):
                info_log = gr.Textbox(
                    label="",
                    placeholder="Generation info will appear here...",
                    lines=12,
                    max_lines=20,
                    interactive=False,
                    elem_classes="info-log"
                )
        
        # Right column - Output
        with gr.Column(scale=1, min_width=320):
            video_output = gr.Video(
                label="🎥 Generated Video",
                autoplay=True,
                height=400,
                elem_classes="video-output"
            )
            
            final_prompt_output = gr.Textbox(
                label="📝 Final Prompt Used",
                interactive=False,
                lines=3
            )
            
            gr.Markdown(
                """
                <p style="text-align: center; margin-top: 10px; font-weight: 700; color: #1F2937;">
                    💡 Right-click on the video to save, or use the download button!
                </p>
                """
            )

    # Define inputs and outputs
    inputs = [
        prompt_input,
        negative_prompt_input,
        enhance_prompt_checkbox,
        duration_slider,
        guidance_scale_slider,
        num_inference_steps_slider,
        height_slider,
        width_slider,
        seed_slider,
        randomize_seed_checkbox,
    ]
    
    outputs = [
        video_output,
        seed_slider,
        final_prompt_output,
        info_log,
    ]

    # Generate button click
    generate_btn.click(
        fn=generate_video,
        inputs=inputs,
        outputs=outputs,
    )

    gr.api(generate_video, api_name="generate_video")

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