# -*- coding: utf-8 -*-
import gc
import os
import random
import subprocess
import sys
import uuid
from pathlib import Path

import gradio as gr
import numpy as np
import spaces
import torch
from huggingface_hub import snapshot_download
from moviepy.editor import AudioFileClip, VideoFileClip
from PIL import Image


ROOT = Path(__file__).resolve().parent
SOURCE_DIR = ROOT / "echomimic_v2"
WEIGHTS_DIR = ROOT / "pretrained_weights"
OUTPUT_DIR = ROOT / "outputs"
SOURCE_REPO = "https://github.com/antgroup/echomimic_v2.git"
WEIGHTS_REPO = "BadToBest/EchoMimicV2"
POSE_NAMES = ["01", "02", "03", "04", "fight", "good", "salute", "ultraman"]

os.environ.setdefault("GRADIO_SSR_MODE", "0")
os.environ.setdefault("GRADIO_ANALYTICS_ENABLED", "False")
os.environ.setdefault("FFMPEG_PATH", "/usr/bin")


def ensure_source():
    if not (SOURCE_DIR / "src").is_dir():
        print("[SETUP] Cloning official EchoMimicV2 source...", flush=True)
        subprocess.run(
            ["git", "clone", "--depth", "1", SOURCE_REPO, str(SOURCE_DIR)],
            check=True,
        )
    if str(SOURCE_DIR) not in sys.path:
        sys.path.insert(0, str(SOURCE_DIR))


ensure_source()

from diffusers import AutoencoderKL, DDIMScheduler
from src.models.pose_encoder import PoseEncoder
from src.models.unet_2d_condition import UNet2DConditionModel
from src.models.unet_3d_emo import EMOUNet3DConditionModel
from src.models.whisper.audio2feature import load_audio_model
from src.pipelines.pipeline_echomimicv2_acc import EchoMimicV2Pipeline
from src.utils.dwpose_util import draw_pose_select_v2
from src.utils.util import save_videos_grid


pipeline = None
weight_dtype = torch.float16
device = "cuda"


def ensure_models():
    WEIGHTS_DIR.mkdir(parents=True, exist_ok=True)
    required = [
        "denoising_unet_acc.pth",
        "motion_module_acc.pth",
        "pose_encoder.pth",
        "reference_unet.pth",
    ]
    if not all((WEIGHTS_DIR / name).is_file() for name in required):
        print("[SETUP] Downloading official EchoMimicV2 accelerated weights...", flush=True)
        snapshot_download(
            repo_id=WEIGHTS_REPO,
            local_dir=WEIGHTS_DIR,
            allow_patterns=required,
        )

    vae_dir = WEIGHTS_DIR / "sd-vae-ft-mse"
    if not (vae_dir / "config.json").is_file():
        print("[SETUP] Downloading SD VAE...", flush=True)
        snapshot_download(repo_id="stabilityai/sd-vae-ft-mse", local_dir=vae_dir)

    base_dir = WEIGHTS_DIR / "sd-image-variations-diffusers"
    if not (base_dir / "unet" / "config.json").is_file():
        print("[SETUP] Downloading SD image-variations base model...", flush=True)
        snapshot_download(
            repo_id="lambdalabs/sd-image-variations-diffusers",
            local_dir=base_dir,
        )

    whisper_path = WEIGHTS_DIR / "audio_processor" / "tiny.pt"
    if not whisper_path.is_file():
        print("[SETUP] Downloading Whisper tiny checkpoint...", flush=True)
        whisper_path.parent.mkdir(parents=True, exist_ok=True)
        torch.hub.download_url_to_file(
            "https://openaipublic.azureedge.net/main/whisper/models/"
            "65147644a518d12f04e32d6f3b26facc3f8dd46e5390956a9424a650c0ce22b9/"
            "tiny.pt",
            str(whisper_path),
        )


def log_model(name, model, source):
    first = next(model.parameters(), None) if isinstance(model, torch.nn.Module) else None
    parameters = (
        sum(parameter.numel() for parameter in model.parameters())
        if isinstance(model, torch.nn.Module)
        else None
    )
    print(
        f"[MODEL] name={name}, class={model.__class__.__name__}, source={source}, "
        f"parameters={parameters}, dtype={getattr(first, 'dtype', None)}, "
        f"device={getattr(first, 'device', None)}",
        flush=True,
    )


def load_models():
    global pipeline
    if pipeline is not None:
        return pipeline
    if not torch.cuda.is_available():
        raise RuntimeError("EchoMimicV2 Accelerated requires CUDA")

    ensure_models()
    base_dir = WEIGHTS_DIR / "sd-image-variations-diffusers"

    print("[MODEL LOAD] VAE from_pretrained", flush=True)
    vae = AutoencoderKL.from_pretrained(
        WEIGHTS_DIR / "sd-vae-ft-mse",
        torch_dtype=weight_dtype,
    ).to(device=device, dtype=weight_dtype)
    log_model("vae", vae, WEIGHTS_DIR / "sd-vae-ft-mse")

    print("[MODEL LOAD] Reference UNet from_pretrained", flush=True)
    reference_unet = UNet2DConditionModel.from_pretrained(
        base_dir,
        subfolder="unet",
        use_safetensors=False,
    ).to(device=device, dtype=weight_dtype)
    state = torch.load(WEIGHTS_DIR / "reference_unet.pth", map_location="cpu", weights_only=True)
    reference_unet.load_state_dict(state)
    del state
    log_model("reference-unet", reference_unet, WEIGHTS_DIR / "reference_unet.pth")

    print("[MODEL LOAD] Accelerated denoising UNet from_pretrained_2d", flush=True)
    unet_kwargs = {
        "use_inflated_groupnorm": True,
        "unet_use_cross_frame_attention": False,
        "unet_use_temporal_attention": False,
        "use_motion_module": True,
        "cross_attention_dim": 384,
        "motion_module_resolutions": [1, 2, 4, 8],
        "motion_module_mid_block": True,
        "motion_module_decoder_only": False,
        "motion_module_type": "Vanilla",
        "motion_module_kwargs": {
            "num_attention_heads": 8,
            "num_transformer_block": 1,
            "attention_block_types": ["Temporal_Self", "Temporal_Self"],
            "temporal_position_encoding": True,
            "temporal_position_encoding_max_len": 32,
            "temporal_attention_dim_div": 1,
        },
    }
    denoising_unet = EMOUNet3DConditionModel.from_pretrained_2d(
        str(base_dir),
        str(WEIGHTS_DIR / "motion_module_acc.pth"),
        subfolder="unet",
        unet_additional_kwargs=unet_kwargs,
    ).to(device=device, dtype=weight_dtype)
    state = torch.load(
        WEIGHTS_DIR / "denoising_unet_acc.pth",
        map_location="cpu",
        weights_only=True,
    )
    missing, unexpected = denoising_unet.load_state_dict(state, strict=False)
    del state
    print(
        f"[MODEL] accelerated-unet missing={len(missing)}, unexpected={len(unexpected)}",
        flush=True,
    )
    log_model("accelerated-denoising-unet", denoising_unet, WEIGHTS_DIR / "denoising_unet_acc.pth")

    pose_encoder = PoseEncoder(
        320,
        conditioning_channels=3,
        block_out_channels=(16, 32, 96, 256),
    ).to(device=device, dtype=weight_dtype)
    state = torch.load(WEIGHTS_DIR / "pose_encoder.pth", map_location="cpu", weights_only=True)
    pose_encoder.load_state_dict(state)
    del state
    log_model("pose-encoder", pose_encoder, WEIGHTS_DIR / "pose_encoder.pth")

    print("[MODEL LOAD] Whisper tiny audio guider", flush=True)
    audio_guider = load_audio_model(
        model_path=str(WEIGHTS_DIR / "audio_processor" / "tiny.pt"),
        device=device,
    )
    log_model("audio-guider", audio_guider, WEIGHTS_DIR / "audio_processor" / "tiny.pt")

    scheduler = DDIMScheduler(
        beta_start=0.00085,
        beta_end=0.012,
        beta_schedule="linear",
        clip_sample=False,
        steps_offset=1,
        prediction_type="v_prediction",
        rescale_betas_zero_snr=True,
        timestep_spacing="trailing",
    )
    pipeline = EchoMimicV2Pipeline(
        vae=vae,
        reference_unet=reference_unet,
        denoising_unet=denoising_unet,
        audio_guider=audio_guider,
        pose_encoder=pose_encoder,
        scheduler=scheduler,
    ).to(device, dtype=weight_dtype)
    gc.collect()
    torch.cuda.empty_cache()
    print("[MODEL LOAD COMPLETE] EchoMimicV2 Accelerated is ready", flush=True)
    return pipeline


def estimate_gpu_duration(_image, audio, *_args):
    try:
        clip = AudioFileClip(audio)
        duration = float(clip.duration)
        clip.close()
    except Exception:
        duration = 5.0
    return max(180, min(900, int(150 + duration * 35)))


def resolve_pose_dir(name):
    pose_dir = SOURCE_DIR / "assets" / "halfbody_demo" / "pose" / name
    if not pose_dir.is_dir():
        raise gr.Error(f"姿态模板不存在：{name}")
    frames = sorted(pose_dir.glob("*.npy"), key=lambda path: int(path.stem))
    if not frames:
        raise gr.Error(f"姿态模板没有可用帧：{name}")
    return pose_dir, len(frames)


@spaces.GPU(size="xlarge", duration=estimate_gpu_duration)
@torch.inference_mode()
def generate(image, audio, pose_name, max_frames, seed, progress=gr.Progress()):
    if not image:
        raise gr.Error("请上传参考图片")
    if not audio:
        raise gr.Error("请上传驱动音频")

    progress(0.05, desc="加载 Accelerated 模型...")
    pipe = load_models()
    pose_dir, pose_frames = resolve_pose_dir(pose_name)

    width = height = 768
    fps = 24
    steps = 6
    cfg = 1.0
    context_frames = 12
    context_overlap = 3
    max_frames = max(12, min(120, int(max_frames)))

    audio_clip = AudioFileClip(audio)
    audio_frames = int(audio_clip.duration * fps)
    length = min(max_frames, audio_frames, pose_frames)
    if length < 12:
        audio_clip.close()
        raise gr.Error("音频过短，至少需要约0.5秒")
    audio_clip = audio_clip.set_duration(length / fps)

    actual_seed = random.randint(1, 2**31 - 1) if seed is None or seed < 0 else int(seed)
    generator = torch.Generator(device=device).manual_seed(actual_seed)
    ref_image = Image.open(image).convert("RGB").resize((width, height))

    progress(0.15, desc="准备姿态模板...")
    pose_list = []
    for index in range(length):
        target = np.zeros((width, height, 3), dtype=np.uint8)
        detected_pose = np.load(pose_dir / f"{index}.npy", allow_pickle=True).tolist()
        imh_new, imw_new, rb, re, cb, ce = detected_pose["draw_pose_params"]
        rendered = draw_pose_select_v2(detected_pose, imh_new, imw_new, ref_w=800)
        rendered = np.transpose(np.asarray(rendered), (1, 2, 0))
        target[rb:re, cb:ce, :] = rendered
        pose_list.append(
            torch.from_numpy(target).to(device=device, dtype=weight_dtype).permute(2, 0, 1) / 255.0
        )
    poses = torch.stack(pose_list, dim=1).unsqueeze(0)

    progress(0.25, desc=f"生成 {length} 帧视频...")
    video = pipe(
        ref_image,
        audio,
        poses,
        width,
        height,
        length,
        steps,
        cfg,
        generator=generator,
        audio_sample_rate=16000,
        context_frames=context_frames,
        fps=fps,
        context_overlap=context_overlap,
        start_idx=0,
    ).videos

    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    job_id = uuid.uuid4().hex
    silent_path = OUTPUT_DIR / f"{job_id}_silent.mp4"
    result_path = OUTPUT_DIR / f"{job_id}.mp4"
    save_videos_grid(video[:, :, :length], str(silent_path), n_rows=1, fps=fps)

    progress(0.92, desc="合并音频...")
    video_clip = VideoFileClip(str(silent_path)).set_audio(audio_clip)
    video_clip.write_videofile(
        str(result_path),
        codec="libx264",
        audio_codec="aac",
        fps=fps,
        threads=2,
        logger=None,
    )
    video_clip.close()
    audio_clip.close()
    silent_path.unlink(missing_ok=True)
    progress(1.0, desc="完成")
    return str(result_path), actual_seed, f"生成 {length} 帧（{length / fps:.1f} 秒），姿态模板：{pose_name}"


with gr.Blocks(title="EchoMimicV2 Accelerated") as demo:
    gr.Markdown(
        "# EchoMimicV2 Accelerated\n"
        "上传半身参考图和音频，使用官方加速权重与预置姿态模板生成说话视频。"
    )
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(label="半身参考图片", type="filepath", sources=["upload"])
            audio_input = gr.Audio(label="驱动音频", type="filepath", sources=["upload"])
            pose_input = gr.Dropdown(
                choices=POSE_NAMES,
                value="01",
                label="动作/姿态模板",
            )
            max_frames_input = gr.Slider(
                minimum=24,
                maximum=120,
                value=72,
                step=12,
                label="最大生成帧数（24fps）",
            )
            seed_input = gr.Number(value=-1, precision=0, label="随机种子（-1为随机）")
            generate_button = gr.Button("生成视频", variant="primary")
        with gr.Column():
            video_output = gr.Video(label="生成结果", autoplay=False)
            seed_output = gr.Number(label="实际种子", precision=0)
            info_output = gr.Textbox(label="生成信息")

    gr.Markdown(
        "首次运行需下载约17GB模型。官方加速配置为768×768、6步；"
        "姿态模板决定身体动作，音频主要控制口型和节奏。"
    )
    generate_button.click(
        generate,
        inputs=[image_input, audio_input, pose_input, max_frames_input, seed_input],
        outputs=[video_output, seed_output, info_output],
    )

demo.queue(default_concurrency_limit=1, max_size=6)

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