import os
import spaces
import torch
from diffusers import WanPipeline
from diffusers.utils.export_utils import export_to_video
import gradio as gr
import tempfile
import numpy as np
import random
import gc


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

# =========================================================
# MODEL CONFIGURATION - Using lighter 1.3B model
# =========================================================
MODEL_ID = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"

LANDSCAPE_WIDTH = 832
LANDSCAPE_HEIGHT = 480
MAX_SEED = np.iinfo(np.int32).max

FIXED_FPS = 16
MIN_FRAMES_MODEL = 9  # Must be 4k+1: 5, 9, 13, 17...
MAX_FRAMES_MODEL = 81

MIN_DURATION = round(MIN_FRAMES_MODEL/FIXED_FPS, 1)
MAX_DURATION = round(MAX_FRAMES_MODEL/FIXED_FPS, 1)

# =========================================================
# MODEL LOADING
# =========================================================
print("Loading Wan 1.3B pipeline...")
pipe = WanPipeline.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
).to('cuda')

# Memory cleanup
for i in range(3): 
    gc.collect()
    torch.cuda.synchronize() 
    torch.cuda.empty_cache()

print("Pipeline loaded successfully!")

# =========================================================
# DEFAULT PROMPTS
# =========================================================
default_prompt_t2v = "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage."
default_negative_prompt = "low quality, worst quality, blurry, distorted, deformed, ugly, bad anatomy, static, watermark, text, subtitle, oversaturated, underexposed"

# =========================================================
# PROMPT ENHANCEMENT
# =========================================================
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
"""


def enhance_prompt(prompt: str) -> str:
    """Enhance the user prompt using Groq LLM API."""
    if not GROQ_API_KEY:
        return prompt
    
    try:
        from groq import Groq
        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


# =========================================================
# GENERATION FUNCTIONS
# =========================================================
@spaces.GPU(duration=300)
def generate_video(
    prompt,
    negative_prompt=default_negative_prompt,
    enhance_prompt_option=False,
    duration_seconds=MAX_DURATION,
    guidance_scale=5.0,
    steps=30,
    seed=42,
    randomize_seed=False,
    progress=gr.Progress(track_tqdm=True),
):
    # Clear memory before generation
    gc.collect()
    torch.cuda.empty_cache()
    
    # Enhance prompt if option is enabled
    final_prompt = prompt
    if enhance_prompt_option:
        final_prompt = enhance_prompt(prompt)
        print(f"Enhanced Prompt: {final_prompt}")
    
    # Calculate num_frames - must satisfy (num_frames - 1) % 4 == 0
    raw_frames = int(round(duration_seconds * FIXED_FPS))
    k = round((raw_frames - 1) / 4)
    num_frames = 4 * k + 1
    num_frames = np.clip(num_frames, MIN_FRAMES_MODEL, MAX_FRAMES_MODEL)
    
    current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)

    output_frames_list = pipe(
        prompt=final_prompt,
        negative_prompt=negative_prompt,
        height=LANDSCAPE_HEIGHT,
        width=LANDSCAPE_WIDTH,
        num_frames=num_frames,
        guidance_scale=float(guidance_scale),
        num_inference_steps=int(steps),
        generator=torch.Generator(device="cuda").manual_seed(current_seed),
    ).frames[0]

    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmpfile:
        video_path = tmpfile.name

    export_to_video(output_frames_list, video_path, fps=FIXED_FPS)
    
    # Clear memory after generation
    gc.collect()
    torch.cuda.empty_cache()

    # Build info log
    actual_duration = num_frames / FIXED_FPS
    info_log = f"""✅ VIDEO GENERATION COMPLETE!
{'=' * 50}
🎬 Video Info:
   • Duration: {actual_duration:.2f} seconds
   • Total Frames: {num_frames}
   • FPS: {FIXED_FPS}
   • Resolution: {LANDSCAPE_WIDTH} x {LANDSCAPE_HEIGHT}
{'=' * 50}
⚙️ Generation Settings:
   • Guidance Scale: {guidance_scale}
   • Inference Steps: {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;
}

/* ===== 🎨 라벨 스타일 ===== */
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;
}

/* ===== 🎨 프로그레스 바 ===== */
.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;
}

/* ===== 반응형 조정 ===== */
@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:
    gr.LoginButton(value="Option: HuggingFace 'Login' for extra GPU quota +", size="sm")       
    # 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_t2v,
                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_seconds_input = gr.Slider(
                minimum=MIN_DURATION,
                maximum=MAX_DURATION,
                step=0.1,
                value=2.0,
                label="⏱️ Duration (seconds)",
                info=f"Range: {MIN_DURATION}s - {MAX_DURATION}s at {FIXED_FPS}fps"
            )
            
            generate_button = 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=3
                )
                seed_input = gr.Slider(
                    label="Seed",
                    minimum=0,
                    maximum=MAX_SEED,
                    step=1,
                    value=42,
                    interactive=True
                )
                randomize_seed_checkbox = gr.Checkbox(
                    label="Randomize Seed",
                    value=True,
                    interactive=True
                )
                steps_slider = gr.Slider(
                    minimum=10,
                    maximum=50,
                    step=1,
                    value=30,
                    label="Inference Steps"
                )
                guidance_scale_input = gr.Slider(
                    minimum=1.0,
                    maximum=15.0,
                    step=0.5,
                    value=5.0,
                    label="Guidance Scale"
                )
            
            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,
                interactive=False,
                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>
                """
            )
    
    # Examples section
    gr.Examples(
        examples=[
            ["POV selfie video, white cat with sunglasses standing on surfboard, relaxed smile, tropical beach behind."],
            ["Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage."],
            ["A cinematic shot of a boat sailing on a calm sea at sunset."],
            ["Drone footage flying over a futuristic city with flying cars."],
        ],
        inputs=[prompt_input],
        cache_examples=False,
        elem_id="examples"
    )
    
    # Define inputs and outputs
    ui_inputs = [
        prompt_input,
        negative_prompt_input,
        enhance_prompt_checkbox,
        duration_seconds_input,
        guidance_scale_input,
        steps_slider,
        seed_input,
        randomize_seed_checkbox
    ]
    
    ui_outputs = [
        video_output,
        seed_input,
        final_prompt_output,
        info_log
    ]
    
    generate_button.click(
        fn=generate_video,
        inputs=ui_inputs,
        outputs=ui_outputs
    )

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