import gc
import random
import tempfile
from functools import lru_cache
from typing import List, Optional, Tuple

import gradio as gr
import numpy as np
import torch
from PIL import Image
from briarmbg import BriaRMBG
from diffusers import AutoencoderKLWan, UniPCMultistepScheduler, WanVACEPipeline
from diffusers.utils import export_to_video

MODEL_ID = "Wan-AI/Wan2.1-VACE-14B-diffusers"
RMBG_ID = "briaai/RMBG-1.4"
LORA_REPO = "vrgamedevgirl84/Wan14BT2VFusioniX"
LORA_FILE = "FusionX_LoRa/Phantom_Wan_14B_FusionX_LoRA.safetensors"
LORA_ADAPTER = "phantom"

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
PIPE_DTYPE = torch.bfloat16 if DEVICE == "cuda" else torch.float32
VAE_DTYPE = torch.float32
RMBG_DTYPE = torch.float32

MOD_VALUE = 32
DEFAULT_H_SLIDER_VALUE = 512
DEFAULT_W_SLIDER_VALUE = 896
NEW_FORMULA_MAX_AREA = 480.0 * 832.0
SLIDER_MIN_H, SLIDER_MAX_H = 128, 896
SLIDER_MIN_W, SLIDER_MAX_W = 128, 896
MAX_SEED = np.iinfo(np.int32).max
FIXED_FPS = 16
MIN_FRAMES_MODEL = 8
MAX_FRAMES_MODEL = 81

MODE_PROMPTS = {
    "Reference": "the playful penguin picks up the green cat eye sunglasses and puts them on",
    "First - Last Frame": "CG animation style, a small blue bird takes off from the ground, flapping its wings. The bird's feathers are delicate, with a unique pattern on its chest. The background shows a blue sky with white clouds under bright sunshine. The camera follows the bird upward, capturing its flight and the vastness of the sky from a close-up, low-angle perspective.",
    "Random Transitions": "Various different characters appear and disappear in a fast transition video showcasing their unique features and personalities. The video is about showcasing different dance styles, with each character performing a distinct dance move. The background is a vibrant, colorful stage with dynamic lighting that changes with each dance style. The camera captures close-ups of the dancers' expressions and movements. Highly dynamic, fast-paced music video, with quick cuts and transitions between characters, cinematic, vibrant colors",
}

DEFAULT_NEGATIVE_PROMPT = (
    "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, "
    "static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, "
    "poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, "
    "messy background, three legs, many people in the background, walking backwards, watermark, text, signature"
)


def clear_cuda():
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        if hasattr(torch.cuda, "ipc_collect"):
            torch.cuda.ipc_collect()


@lru_cache(maxsize=1)
def get_pipe() -> WanVACEPipeline:
    vae = AutoencoderKLWan.from_pretrained(MODEL_ID, subfolder="vae", torch_dtype=VAE_DTYPE)
    pipe = WanVACEPipeline.from_pretrained(MODEL_ID, vae=vae, torch_dtype=PIPE_DTYPE)
    pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=2.0)
    pipe.load_lora_weights(LORA_REPO, weight_name=LORA_FILE, adapter_name=LORA_ADAPTER)
    try:
        pipe.set_adapters([LORA_ADAPTER], adapter_weights=[1.0])
    except Exception:
        pass
    pipe = pipe.to(DEVICE)
    return pipe


@lru_cache(maxsize=1)
def get_rmbg() -> BriaRMBG:
    model = BriaRMBG.from_pretrained(RMBG_ID)
    model = model.to(DEVICE, dtype=RMBG_DTYPE)
    model.eval()
    return model


def remove_alpha_channel(image: Image.Image) -> Image.Image:
    if image.mode in ("RGBA", "LA"):
        background = Image.new("RGB", image.size, (255, 255, 255))
        paste_source = image if image.mode == "RGBA" else image.convert("RGB")
        background.paste(paste_source, mask=image.split()[-1])
        return background
    if image.mode == "P":
        if "transparency" in image.info:
            image = image.convert("RGBA")
            background = Image.new("RGB", image.size, (255, 255, 255))
            background.paste(image, mask=image.split()[-1])
            return background
        return image.convert("RGB")
    if image.mode != "RGB":
        return image.convert("RGB")
    return image


def numpy2pytorch(imgs: List[np.ndarray]) -> torch.Tensor:
    h = torch.from_numpy(np.stack(imgs, axis=0)).float() / 127.0 - 1.0
    return h.movedim(-1, 1)


def resize_without_crop(image: np.ndarray, target_width: int, target_height: int) -> np.ndarray:
    return np.array(Image.fromarray(image).resize((target_width, target_height), Image.LANCZOS))


def run_rmbg(img: np.ndarray, sigma: float = 0.0) -> Tuple[np.ndarray, np.ndarray]:
    model = get_rmbg()
    h, w, c = img.shape
    if c != 3:
        raise ValueError("RMBG expects RGB input")
    k = (256.0 / float(h * w)) ** 0.5
    feed = resize_without_crop(img, int(64 * round(w * k)), int(64 * round(h * k)))
    feed = numpy2pytorch([feed]).to(device=DEVICE, dtype=RMBG_DTYPE)
    with torch.inference_mode():
        alpha = model(feed)[0][0]
    alpha = torch.nn.functional.interpolate(alpha, size=(h, w), mode="bilinear")
    alpha = alpha.movedim(1, -1)[0].detach().float().cpu().numpy().clip(0, 1)
    result = 127 + (img.astype(np.float32) - 127 + sigma) * alpha
    return result.clip(0, 255).astype(np.uint8), alpha


def remove_background_from_image(image: Image.Image) -> Image.Image:
    image = remove_alpha_channel(image)
    result_array, alpha_mask = run_rmbg(np.array(image))
    result_image = Image.fromarray(result_array).convert("RGBA")
    alpha_mask_2d = np.squeeze(alpha_mask)
    if alpha_mask_2d.ndim > 2:
        alpha_mask_2d = alpha_mask_2d[:, :, 0]
    alpha_array = (alpha_mask_2d * 255).astype(np.uint8)
    result_image.putalpha(Image.fromarray(alpha_array, "L"))
    return result_image


def calculate_new_dimensions_wan(
    pil_image: Image.Image,
    mod_val: int,
    calculation_max_area: float,
    min_slider_h: int,
    max_slider_h: int,
    min_slider_w: int,
    max_slider_w: int,
    default_h: int,
    default_w: int,
) -> Tuple[int, int]:
    orig_w, orig_h = pil_image.size
    if orig_w <= 0 or orig_h <= 0:
        return default_h, default_w
    aspect_ratio = orig_h / orig_w
    calc_h = round(np.sqrt(calculation_max_area * aspect_ratio))
    calc_w = round(np.sqrt(calculation_max_area / aspect_ratio))
    calc_h = max(mod_val, (calc_h // mod_val) * mod_val)
    calc_w = max(mod_val, (calc_w // mod_val) * mod_val)
    new_h = int(np.clip(calc_h, min_slider_h, (max_slider_h // mod_val) * mod_val))
    new_w = int(np.clip(calc_w, min_slider_w, (max_slider_w // mod_val) * mod_val))
    return new_h, new_w


def handle_gallery_upload_for_dims_wan(gallery_images, current_h_val, current_w_val):
    if not gallery_images:
        return gr.update(value=DEFAULT_H_SLIDER_VALUE), gr.update(value=DEFAULT_W_SLIDER_VALUE)
    try:
        first_image = gallery_images[0][0] if isinstance(gallery_images[0], (list, tuple)) else gallery_images[0]
        first_image = remove_alpha_channel(first_image)
        new_h, new_w = calculate_new_dimensions_wan(
            first_image,
            MOD_VALUE,
            NEW_FORMULA_MAX_AREA,
            SLIDER_MIN_H,
            SLIDER_MAX_H,
            SLIDER_MIN_W,
            SLIDER_MAX_W,
            DEFAULT_H_SLIDER_VALUE,
            DEFAULT_W_SLIDER_VALUE,
        )
        return gr.update(value=new_h), gr.update(value=new_w)
    except Exception:
        gr.Warning("Error attempting to calculate new dimensions")
        return gr.update(value=DEFAULT_H_SLIDER_VALUE), gr.update(value=DEFAULT_W_SLIDER_VALUE)


def update_prompt_from_mode(mode: str) -> str:
    return MODE_PROMPTS.get(mode, "")


def prepare_video_and_mask_ref2v(height: int, width: int, num_frames: int):
    frames = [Image.new("RGB", (width, height), (128, 128, 128)) for _ in range(num_frames)]
    mask_white = Image.new("L", (width, height), 255)
    mask = [mask_white for _ in range(num_frames)]
    return frames, mask


def prepare_video_and_mask_flf2v(first_img: Image.Image, last_img: Image.Image, height: int, width: int, num_frames: int):
    first_img = remove_alpha_channel(first_img).resize((width, height))
    last_img = remove_alpha_channel(last_img).resize((width, height))
    frames = [first_img]
    frames.extend([Image.new("RGB", (width, height), (128, 128, 128)) for _ in range(num_frames - 2)])
    frames.append(last_img)
    mask_black = Image.new("L", (width, height), 0)
    mask_white = Image.new("L", (width, height), 255)
    mask = [mask_black] + [mask_white for _ in range(num_frames - 2)] + [mask_black]
    return frames, mask


def calculate_random2v_frame_indices(num_images: int, num_frames: int) -> List[int]:
    if num_images <= 0:
        return []
    if num_images == 1:
        return [num_frames // 2]
    if num_images >= num_frames:
        return list(range(num_frames))
    indices = []
    step = (num_frames - 1) / (num_images - 1)
    for i in range(num_images):
        frame_idx = min(int(round(i * step)), num_frames - 1)
        indices.append(frame_idx)
    seen = set()
    unique_indices = []
    for idx in indices:
        if idx not in seen:
            seen.add(idx)
            unique_indices.append(idx)
    return unique_indices


def prepare_video_and_mask_random2v(images: List[Image.Image], frame_indices: List[int], height: int, width: int, num_frames: int):
    clean_images = [remove_alpha_channel(img).resize((width, height)) for img in images]
    frames = [Image.new("RGB", (width, height), (128, 128, 128)) for _ in range(num_frames)]
    mask_black = Image.new("L", (width, height), 0)
    mask_white = Image.new("L", (width, height), 255)
    mask = [mask_white for _ in range(num_frames)]
    for img, idx in zip(clean_images, frame_indices):
        if idx >= num_frames:
            raise ValueError(f"Frame index {idx} exceeds num_frames {num_frames}")
        frames[idx] = img
        mask[idx] = mask_black
    return frames, mask


def normalize_gallery_images(gallery_images) -> List[Image.Image]:
    images: List[Image.Image] = []
    for item in gallery_images or []:
        if isinstance(item, (list, tuple)) and item:
            image = item[0]
        else:
            image = item
        if not isinstance(image, Image.Image):
            raise gr.Error("Gallery contains an invalid image item.")
        images.append(image)
    return images


def generate_video(
    gallery_images,
    mode,
    prompt,
    height,
    width,
    negative_prompt=DEFAULT_NEGATIVE_PROMPT,
    duration_seconds=2.0,
    guidance_scale=1.0,
    steps=6,
    seed=42,
    randomize_seed=False,
    remove_bg=False,
    progress=gr.Progress(track_tqdm=True),
):
    images = normalize_gallery_images(gallery_images)
    if not images:
        raise gr.Error("Please upload at least one image to the gallery.")

    processed_images: List[Image.Image] = []
    for image in images:
        if mode == "Reference" and remove_bg:
            image = remove_background_from_image(image)
        image = remove_alpha_channel(image)
        processed_images.append(image)

    if mode == "First - Last Frame":
        if len(processed_images) < 2:
            raise gr.Error(f"First - Last Frame mode requires at least 2 images, but only {len(processed_images)} were supplied.")
        processed_images = processed_images[:2]

    target_h = max(MOD_VALUE, (int(height) // MOD_VALUE) * MOD_VALUE)
    target_w = max(MOD_VALUE, (int(width) // MOD_VALUE) * MOD_VALUE)
    num_frames = int(np.clip(int(round(duration_seconds * FIXED_FPS)), MIN_FRAMES_MODEL, MAX_FRAMES_MODEL))
    current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)

    if mode == "First - Last Frame":
        frames, mask = prepare_video_and_mask_flf2v(processed_images[0], processed_images[1], target_h, target_w, num_frames)
        reference_images = None
    elif mode == "Reference":
        frames, mask = prepare_video_and_mask_ref2v(target_h, target_w, num_frames)
        reference_images = processed_images
    else:
        frame_indices = calculate_random2v_frame_indices(len(processed_images), num_frames)
        frames, mask = prepare_video_and_mask_random2v(processed_images, frame_indices, target_h, target_w, num_frames)
        reference_images = None

    pipe = get_pipe()
    clear_cuda()

    generator_device = DEVICE if DEVICE == "cuda" else "cpu"
    generator = torch.Generator(device=generator_device).manual_seed(current_seed)

    with torch.inference_mode():
        output_frames_list = pipe(
            video=frames,
            mask=mask,
            reference_images=reference_images,
            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=".mp4", delete=False) as tmpfile:
        video_path = tmpfile.name
        export_to_video(output_frames_list, video_path, fps=FIXED_FPS)

    clear_cuda()
    return video_path, current_seed


control_modes = """
**3 Control Modes Available:**
- **Reference** Generate a video incorporating elements from input reference images
- **First - Last Frame** Generate a video using first and last frame conditioning defined by input images
- **Random Transitions** Generate a video with intermediate transitions between multiple input images
"""


with gr.Blocks() as demo:
    gr.Markdown("# Fast 6 step Wan 2.1 VACE (14B)")
    gr.Markdown(
        "Using [**Wan2.1-VACE-14B**](https://huggingface.co/Wan-AI/Wan2.1-VACE-14B-diffusers) + "
        "[**FusionX Phantom LoRA**](https://huggingface.co/vrgamedevgirl84/Wan14BT2VFusioniX) for fast video generation with multiple conditions."
    )
    gr.Markdown(control_modes)

    with gr.Row():
        with gr.Column():
            with gr.Group():
                mode_radio = gr.Radio(
                    choices=["Reference", "First - Last Frame", "Random Transitions"],
                    value="Reference",
                    label="Control Mode",
                )
                gallery_component = gr.Gallery(
                    label="Upload 1 or more images",
                    show_label=True,
                    elem_id="gallery",
                    columns=3,
                    rows=2,
                    object_fit="contain",
                    height="auto",
                    type="pil",
                    allow_preview=True,
                )
                remove_bg_checkbox = gr.Checkbox(
                    label="Remove Background",
                    value=False,
                    visible=True,
                    info="Removes background from input images. Best used in Reference mode.",
                )
                prompt_input = gr.Textbox(label="Prompt", value=MODE_PROMPTS["Reference"])
                duration_seconds_input = gr.Slider(
                    minimum=round(MIN_FRAMES_MODEL / FIXED_FPS, 1),
                    maximum=round(MAX_FRAMES_MODEL / FIXED_FPS, 1),
                    step=0.1,
                    value=2.3,
                    label="Duration (seconds)",
                    info=f"Clamped to model's {MIN_FRAMES_MODEL}-{MAX_FRAMES_MODEL} frames at {FIXED_FPS} fps.",
                )
                with gr.Accordion("Advanced Settings", 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)
                    with gr.Row():
                        height_input = gr.Slider(minimum=SLIDER_MIN_H, maximum=SLIDER_MAX_H, step=MOD_VALUE, value=DEFAULT_H_SLIDER_VALUE, label=f"Output Height (multiple of {MOD_VALUE})")
                        width_input = gr.Slider(minimum=SLIDER_MIN_W, maximum=SLIDER_MAX_W, step=MOD_VALUE, value=DEFAULT_W_SLIDER_VALUE, label=f"Output Width (multiple of {MOD_VALUE})")
                    steps_slider = gr.Slider(minimum=1, maximum=10, step=1, value=6, label="Inference Steps")
                    guidance_scale_input = gr.Slider(minimum=0.0, maximum=5.0, step=0.5, value=1.0, label="Guidance Scale", visible=False)
                generate_button = gr.Button("Generate Video", variant="primary")

        with gr.Column():
            video_output = gr.Video(label="Generated Video", autoplay=True, interactive=False)

    def update_bg_removal_visibility(mode: str):
        return gr.update(visible=(mode == "Reference"))

    mode_radio.change(fn=update_prompt_from_mode, inputs=[mode_radio], outputs=[prompt_input])
    mode_radio.change(fn=update_bg_removal_visibility, inputs=[mode_radio], outputs=[remove_bg_checkbox])
    gallery_component.change(
        fn=handle_gallery_upload_for_dims_wan,
        inputs=[gallery_component, height_input, width_input],
        outputs=[height_input, width_input],
    )

    ui_inputs = [
        gallery_component,
        mode_radio,
        prompt_input,
        height_input,
        width_input,
        negative_prompt_input,
        duration_seconds_input,
        guidance_scale_input,
        steps_slider,
        seed_input,
        randomize_seed_checkbox,
        remove_bg_checkbox,
    ]

    generate_button.click(fn=generate_video, inputs=ui_inputs, outputs=[video_output, seed_input])

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