import os
import sys
import subprocess
import traceback
import gc
import tempfile
import random
import time
from pathlib import Path

os.system("pip install spaces-0.1.0-py3-none-any.whl moviepy==1.0.3 imageio[ffmpeg] librosa soundfile accelerate")
os.system("pip install git+https://github.com/tolgacangoz/diffusers.git")

import spaces
import torch
import numpy as np
import librosa
import soundfile as sf
from PIL import Image
from moviepy.editor import VideoFileClip, concatenate_videoclips
from huggingface_hub import snapshot_download
import gradio as gr

try:
    import diffusers
    from diffusers import AutoencoderKLWan, WanPipeline, WanImageToVideoPipeline, UniPCMultistepScheduler, WanSpeechToVideoPipeline
    from diffusers.utils import export_to_video
except ImportError:
    pass

MODEL_ID_TI2V = "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers"
MODEL_ID_S2V = "tolgacangoz/Wan2.2-S2V-14B-Diffusers"

MODELS = {
    "ti2v_text": None,
    "ti2v_image": None,
    "s2v": None
}

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

def load_models_at_startup():
    global MODELS
    
    try:
        vae = AutoencoderKLWan.from_pretrained(MODEL_ID_TI2V, subfolder="vae", torch_dtype=torch.float32)
        
        text_pipe = WanPipeline.from_pretrained(MODEL_ID_TI2V, vae=vae, torch_dtype=torch.bfloat16)
        text_pipe.scheduler = UniPCMultistepScheduler.from_config(text_pipe.scheduler.config, flow_shift=8.0)
        
        try:
            if DEVICE == "cuda":
                text_pipe.enable_model_cpu_offload()
            else:
                text_pipe.to(DEVICE)
        except RuntimeError:
            text_pipe.to("cpu")
            
        MODELS["ti2v_text"] = text_pipe

        image_pipe = WanImageToVideoPipeline.from_pretrained(MODEL_ID_TI2V, vae=vae, torch_dtype=torch.bfloat16)
        image_pipe.scheduler = UniPCMultistepScheduler.from_config(image_pipe.scheduler.config, flow_shift=8.0)
        
        try:
            if DEVICE == "cuda":
                image_pipe.enable_model_cpu_offload()
            else:
                image_pipe.to(DEVICE)
        except RuntimeError:
            image_pipe.to("cpu")
            
        MODELS["ti2v_image"] = image_pipe
        
    except Exception as e:
        pass

    try:
        s2v_pipe = WanSpeechToVideoPipeline.from_pretrained(
            MODEL_ID_S2V,
            torch_dtype=torch.bfloat16 
        )
        try:
            if DEVICE == "cuda":
                s2v_pipe.enable_model_cpu_offload()
            else:
                s2v_pipe.to(DEVICE)
        except RuntimeError:
            s2v_pipe.to("cpu")
            
        MODELS["s2v"] = s2v_pipe
    except Exception as e:
        pass

load_models_at_startup()

def auto_duration_estimator(mode, input_data, duration_val):
    base_overhead = 45
    if mode == "s2v":
        audio_path = input_data
        if audio_path:
            try:
                dur = librosa.get_duration(filename=audio_path)
                return int(base_overhead + (dur * 15))
            except:
                return 120
        return 120
    else:
        num_images = len(input_data) if input_data else 0
        if num_images > 0:
            total_seconds = max(duration_val, num_images * 2)
        else:
            total_seconds = duration_val
        return int(base_overhead + (total_seconds * 12))

def fast_stitch_videos(video_paths):
    if not video_paths: return None
    if len(video_paths) == 1: return video_paths[0]
    
    try:
        with tempfile.NamedTemporaryFile(mode='w', suffix='.txt', delete=False) as f:
            for path in video_paths:
                f.write(f"file '{path}'\n")
            list_path = f.name
            
        with tempfile.NamedTemporaryFile(suffix="_stitched_stream.mp4", delete=False) as tmp:
            out_path = tmp.name

        cmd = [
            "ffmpeg", "-y", "-f", "concat", "-safe", "0",
            "-i", list_path, "-c", "copy", out_path
        ]
        subprocess.run(cmd, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        os.remove(list_path)
        return out_path
    except:
        return video_paths[-1]

@spaces(duration=lambda *args: auto_duration_estimator("ti2v", args[0], args[5]))
def generate_ti2v_gpu_stream(input_files, prompt, height, width, negative_prompt, duration_seconds, guidance_scale, steps, seed, randomize_seed, progress=gr.Progress(track_tqdm=True)):
    global MODELS
    text_to_video_pipe = MODELS.get("ti2v_text")
    image_to_video_pipe = MODELS.get("ti2v_image")
    
    if not text_to_video_pipe or not image_to_video_pipe:
        raise gr.Error("Models failed to load at startup.")

    MOD_VALUE = 32
    target_h = max(MOD_VALUE, (int(height) // MOD_VALUE) * MOD_VALUE)
    target_w = max(MOD_VALUE, (int(width) // MOD_VALUE) * MOD_VALUE)
    
    master_seed = random.randint(0, 2**32 - 1) if randomize_seed else int(seed)
    
    video_clips_paths = []
    pil_images = []

    if input_files:
        files_list = input_files if isinstance(input_files, list) else [input_files]
        for f in files_list:
            try:
                path = f.name if hasattr(f, "name") else f
                img = Image.open(path).convert("RGB")
                pil_images.append(img)
            except:
                continue

    SAFE_CHUNK_DURATION = 4.0
    FIXED_FPS = 24
    
    last_preview_frame = None

    if len(pil_images) > 0:
        seconds_per_image = max(2.0, duration_seconds / len(pil_images))
        
        for i, img in enumerate(pil_images):
            current_chunk_duration = min(seconds_per_image, SAFE_CHUNK_DURATION)
            num_frames = int(current_chunk_duration * FIXED_FPS)
            
            local_seed = master_seed + i
            generator = torch.Generator(device=DEVICE).manual_seed(local_seed)
            resized_image = img.resize((target_w, target_h))
            
            try:
                with torch.inference_mode():
                    output_frames = image_to_video_pipe(
                        image=resized_image, 
                        prompt=prompt, 
                        negative_prompt=negative_prompt,
                        height=target_h, 
                        width=target_w, 
                        num_frames=num_frames,
                        guidance_scale=float(guidance_scale), 
                        num_inference_steps=int(steps),
                        generator=generator
                    ).frames[0]
                
                with tempfile.NamedTemporaryFile(suffix=f"_img_{i}.mp4", delete=False) as tmp:
                    export_to_video(output_frames, tmp.name, fps=FIXED_FPS)
                    video_clips_paths.append(tmp.name)
                
                if len(output_frames) > 0:
                    last_preview_frame = output_frames[-1]
                
                current_stitched = fast_stitch_videos(video_clips_paths)
                yield current_stitched, last_preview_frame, master_seed
                
            except Exception:
                continue
    else:
        num_chunks = int(np.ceil(duration_seconds / SAFE_CHUNK_DURATION))
        frames_per_chunk = int(SAFE_CHUNK_DURATION * FIXED_FPS)
        
        for i in range(num_chunks):
            chunk_seed = master_seed + (i * 100)
            generator = torch.Generator(device=DEVICE).manual_seed(chunk_seed)
            
            with torch.inference_mode():
                output_frames = text_to_video_pipe(
                    prompt=prompt, 
                    negative_prompt=negative_prompt,
                    height=target_h, 
                    width=target_w, 
                    num_frames=frames_per_chunk,
                    guidance_scale=float(guidance_scale), 
                    num_inference_steps=int(steps),
                    generator=generator
                ).frames[0]
            
            with tempfile.NamedTemporaryFile(suffix=f"_chunk_{i}.mp4", delete=False) as tmp:
                export_to_video(output_frames, tmp.name, fps=FIXED_FPS)
                video_clips_paths.append(tmp.name)
            
            if len(output_frames) > 0:
                last_preview_frame = output_frames[-1]

            current_stitched = fast_stitch_videos(video_clips_paths)
            yield current_stitched, last_preview_frame, master_seed

def merge_audio_video(video_path, audio_path, output_path):
    cmd = [
        "ffmpeg", "-y",
        "-i", video_path,
        "-i", audio_path,
        "-c:v", "copy",
        "-c:a", "aac",
        "-map", "0:v:0", "-map", "1:a:0",
        "-shortest",
        output_path
    ]
    subprocess.run(cmd, check=True)
    return output_path

def load_audio_for_model(audio_filepath):
    try:
        wav, sr = librosa.load(audio_filepath, sr=16000)
        return wav, sr
    except:
        return None, None

@spaces(duration=lambda *args: auto_duration_estimator("s2v", args[1], 0))
def generate_s2v_gpu(image_input, audio_filepath, prompt, seed, randomize_seed):
    global MODELS
    pipe = MODELS.get("s2v")
    if not pipe:
        raise gr.Error("S2V Model not initialized.")

    if image_input is None or audio_filepath is None:
        raise gr.Error("Inputs Missing")
    
    audio_values, sample_rate = load_audio_for_model(audio_filepath)
    if audio_values is None:
        raise gr.Error("Invalid Audio")

    init_image = image_input.convert("RGB")
    w, h = init_image.size
    w = (w // 16) * 16
    h = (h // 16) * 16
    init_image = init_image.resize((w, h), Image.LANCZOS)
    
    current_seed = random.randint(0, 2**32 - 1) if randomize_seed else int(seed)
    generator = torch.Generator(device=DEVICE).manual_seed(current_seed)

    with torch.inference_mode():
        out = pipe(
            image=init_image,
            audio=audio_values,
            num_inference_steps=25,
            guidance_scale=4.0,
            sampling_rate=sample_rate,
            prompt=prompt,
            generator=generator
        )
    
    frames = out.frames[0]
    
    with tempfile.NamedTemporaryFile(suffix="_temp_mute.mp4", delete=False) as tmp_vid:
        temp_mute_path = tmp_vid.name
    
    with tempfile.NamedTemporaryFile(suffix="_output_s2v.mp4", delete=False) as tmp_final:
        final_video_path = tmp_final.name

    export_to_video(frames, temp_mute_path, fps=30) 
    final_output = merge_audio_video(temp_mute_path, audio_filepath, final_video_path)
    
    return final_output, current_seed

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# Wan 2.2 Unified Streaming Video Platform")
    
    with gr.Tabs():
        with gr.TabItem("Text & Image to Video (Streaming & Long Duration)"):
            with gr.Row():
                with gr.Column(scale=1):
                    ti2v_files = gr.File(label="Input Images", file_count="multiple", type="filepath", file_types=["image"])
                    ti2v_prompt = gr.Textbox(label="Prompt", value="Cinematic view, realistic lighting, 4k", lines=2)
                    ti2v_duration = gr.Slider(minimum=2, maximum=300, step=1, value=5, label="Total Duration (s)")
                    
                    with gr.Accordion("Advanced", open=False):
                        ti2v_neg = gr.Textbox(label="Negative Prompt", value="low quality, distortion, text, watermark", lines=2)
                        ti2v_seed = gr.Slider(label="Seed", minimum=0, maximum=2**32-1, step=1, value=42)
                        ti2v_rand = gr.Checkbox(label="Random Seed", value=True)
                        with gr.Row():
                            ti2v_h = gr.Slider(256, 1024, 32, 832, label="Height")
                            ti2v_w = gr.Slider(256, 1024, 32, 832, label="Width")
                        ti2v_steps = gr.Slider(2, 10, 1, 4, label="Steps")
                        ti2v_scale = gr.Slider(1.0, 8.0, 0.1, 5.0, label="CFG")
                    
                    btn_ti2v = gr.Button("Start Streaming Generation", variant="primary")
                
                with gr.Column(scale=2):
                    with gr.Row():
                        out_ti2v = gr.Video(label="Live Video Stream", autoplay=True)
                        out_preview_ti2v = gr.Image(label="Last Frame Preview", interactive=False)
                    out_seed_ti2v = gr.Number(label="Seed Used")
            
            btn_ti2v.click(
                fn=generate_ti2v_gpu_stream, 
                inputs=[ti2v_files, ti2v_prompt, ti2v_h, ti2v_w, ti2v_neg, ti2v_duration, ti2v_scale, ti2v_steps, ti2v_seed, ti2v_rand], 
                outputs=[out_ti2v, out_preview_ti2v, out_seed_ti2v]
            )

        with gr.TabItem("Speech to Video (S2V)"):
            with gr.Row():
                with gr.Column(scale=1):
                    s2v_img = gr.Image(label="Reference Image", type="pil")
                    s2v_audio = gr.Audio(label="Audio Input", type="filepath")
                    s2v_prompt = gr.Textbox(label="Prompt", value="Realistic movement, talking face")
                    s2v_seed = gr.Slider(label="Seed", minimum=0, maximum=2**32-1, step=1, value=42)
                    s2v_rand = gr.Checkbox(label="Random Seed", value=True)
                    btn_s2v = gr.Button("Generate S2V", variant="primary")
                
                with gr.Column(scale=2):
                    out_s2v = gr.Video(label="Result")
                    out_seed_s2v = gr.Number(label="Seed Used")
            
            btn_s2v.click(generate_s2v_gpu, [s2v_img, s2v_audio, s2v_prompt, s2v_seed, s2v_rand], [out_s2v, out_seed_s2v])

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