"""
AIBRUH/avatar-forcing — AvatarForcing Gradio Space
One-step streaming talking avatar: portrait + audio → MP4
Built on: KlingAIResearch/AvatarForcing + Wan2.1-T2V-1.3B
"""
import os
import sys
import math
import subprocess
import tempfile
import soundfile as sf

import torch
import torchaudio
import gradio as gr
from PIL import Image
from einops import rearrange
from torchvision import transforms
from torchvision.transforms import InterpolationMode
import torchvision.transforms.functional as F_tv
import imageio
from transformers import Wav2Vec2FeatureExtractor, Wav2Vec2Model
from huggingface_hub import snapshot_download
from omegaconf import OmegaConf
from collections import OrderedDict

# ── Paths ─────────────────────────────────────────────────────────────────────
APP_DIR     = os.path.dirname(os.path.abspath(__file__))
REPO_DIR    = os.path.join(APP_DIR, "AvatarForcing")
WAN_DIR     = os.path.join(REPO_DIR, "wan_models", "Wan2.1-T2V-1.3B")
WAV2VEC_DIR = os.path.join(REPO_DIR, "wan_models", "wav2vec2-base-960h")
CKPT_DIR    = os.path.join(REPO_DIR, "checkpoints")

# ── Step 1: Clone repo ────────────────────────────────────────────────────────
if not os.path.exists(os.path.join(REPO_DIR, "inference.py")):
    print("[1/4] Cloning AvatarForcing repo...")
    subprocess.run([
        "git", "clone", "--depth=1",
        "https://github.com/KlingAIResearch/AvatarForcing.git",
        REPO_DIR
    ], check=True)
else:
    print("[1/4] Repo already cloned.")

sys.path.insert(0, REPO_DIR)
os.makedirs(CKPT_DIR, exist_ok=True)
os.makedirs(WAN_DIR, exist_ok=True)
os.makedirs(WAV2VEC_DIR, exist_ok=True)

# ── Step 2: Download models ───────────────────────────────────────────────────
print("[2/4] Downloading models (first boot only)...")

if not os.path.exists(os.path.join(WAN_DIR, "config.json")):
    snapshot_download(
        "Wan-AI/Wan2.1-T2V-1.3B",
        local_dir=WAN_DIR,
        ignore_patterns=["*.msgpack", "flax_model*", "tf_model*", "rust_model*"],
    )

if not os.path.exists(os.path.join(WAV2VEC_DIR, "config.json")):
    snapshot_download("facebook/wav2vec2-base-960h", local_dir=WAV2VEC_DIR)

if not os.path.exists(os.path.join(CKPT_DIR, "model.pt")):
    snapshot_download("lycui/AvatarForcing", local_dir=CKPT_DIR)

print("[2/4] Models ready.")

# ── Step 3: Load pipeline ─────────────────────────────────────────────────────
# MUST chdir to REPO_DIR: wan_wrapper.py uses hardcoded relative paths like
# "./wan_models/Wan2.1-T2V-1.3B/..." so CWD must be the repo root.
os.chdir(REPO_DIR)
print(f"[3/4] CWD → {REPO_DIR}")

# Patch wan_wrapper.py: change text encoder from float32 → bfloat16
# This halves CPU RAM usage for the UMT5-XXL encoder (~10GB → ~5GB)
_ww_path = os.path.join(REPO_DIR, "utils", "wan_wrapper.py")
with open(_ww_path, "r") as _f:
    _ww_src = _f.read()
_ww_patched = _ww_src.replace(
    "dtype=torch.float32,\n            device=torch.device('cpu')",
    "dtype=torch.bfloat16,\n            device=torch.device('cpu')",
)
with open(_ww_path, "w") as _f:
    _f.write(_ww_patched)
print("[3/4] Patched wan_wrapper.py: text encoder dtype float32 → bfloat16")

print("[3/4] Loading AvatarForcing pipeline...")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

from pipeline import AvatarForcingInferencePipeline
from utils.inject import _apply_lora

config_path  = os.path.join(REPO_DIR, "configs", "avatarforcing.yaml")
default_path = os.path.join(REPO_DIR, "configs", "default_config.yaml")
config = OmegaConf.merge(OmegaConf.load(default_path), OmegaConf.load(config_path))

# Override model paths to absolute so CWD doesn't matter
config.data.wav2vec_path = WAV2VEC_DIR
if hasattr(config, "model_kwargs") and hasattr(config.model_kwargs, "model_path"):
    config.model_kwargs.model_path = WAN_DIR

# Build pipeline (loads Wan2.1-T2V-1.3B from WAN_DIR)
pipeline = AvatarForcingInferencePipeline(config, device=device)

# Load AvatarForcing DMD weights
ckpt_path = os.path.join(CKPT_DIR, "model.pt")
state_dict = torch.load(ckpt_path, map_location="cpu")
pipeline.generator.model = _apply_lora(pipeline.generator.model, config["models"]["lora"])
pipeline.generator.load_state_dict(state_dict["generator"])
pipeline = pipeline.to(device=device, dtype=torch.bfloat16)
pipeline.eval()

print("[3/4] Pipeline loaded.")

# ── Step 4: Audio encoder ─────────────────────────────────────────────────────
print("[4/4] Loading Wav2Vec2 audio encoder...")
wav2vec_extractor = Wav2Vec2FeatureExtractor.from_pretrained(WAV2VEC_DIR)
wav2vec_model = Wav2Vec2Model.from_pretrained(WAV2VEC_DIR).eval().to(device)
print("[4/4] Ready.")

# ── Helpers ───────────────────────────────────────────────────────────────────
class ResizeKeepRatioArea16:
    def __init__(self, area_hw=(480, 832), div=16):
        self.A = area_hw[0] * area_hw[1]
        self.d = div

    def __call__(self, img):
        w, h = img.size
        s = min(1.0, math.sqrt(self.A / (h * w)))
        nh = max(self.d, int(h * s) // self.d * self.d)
        nw = max(self.d, int(w * s) // self.d * self.d)
        return F_tv.resize(img, (nh, nw), interpolation=InterpolationMode.BILINEAR, antialias=True)


img_transform = transforms.Compose([
    ResizeKeepRatioArea16((480, 832), 16),
    transforms.ToTensor(),
    transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
])


def encode_audio(wav_path: str, num_frames: int, fps: int = 25) -> torch.Tensor:
    """Replicates TextImageAudioPairDataset audio processing."""
    data, sr = sf.read(wav_path)
    if data.ndim > 1:
        data = data.mean(axis=1)
    data_t = torch.tensor(data, dtype=torch.float32)
    if sr != 16000:
        data_t = torchaudio.functional.resample(data_t, sr, 16000)

    teacher_len = num_frames * 4 + 80
    max_audio_len = int(teacher_len * (16000 / fps))
    data_t = data_t[:max_audio_len]

    inputs = wav2vec_extractor(
        data_t.numpy(), sampling_rate=16000, return_tensors="pt", padding=True
    )
    inputs = {k: v.to(device) for k, v in inputs.items()}
    with torch.no_grad():
        out = wav2vec_model(**inputs, output_hidden_states=True)

    last = out.last_hidden_state  # [1, T, 768]
    # Concatenate hidden states as per dataset.py pattern
    hidden = [h for h in out.hidden_states[1:]] if out.hidden_states else []
    if hidden:
        emb = torch.cat([last] + hidden, dim=-1)
    else:
        emb = last

    # Prepend zero frame (dataset.py padding)
    zero = torch.zeros(1, 1, emb.shape[-1], device=device)
    emb = torch.cat([zero, emb], dim=1)
    return emb


# ── Inference ─────────────────────────────────────────────────────────────────
def generate(portrait_path: str, audio_path: str, prompt: str, num_seconds: int) -> str:
    num_frames = num_seconds * 25 + 1  # 25 fps; must satisfy (frames-1) % 4 == 0
    # Round to nearest valid value
    num_frames = ((num_frames - 1 + 3) // 4) * 4 + 1

    # Image → latent
    img = Image.open(portrait_path).convert("RGB")
    img_t = img_transform(img).unsqueeze(0).unsqueeze(2).to(device=device, dtype=torch.bfloat16)
    initial_latent = pipeline.vae.encode_to_latent(img_t).to(device=device, dtype=torch.bfloat16)

    # Build conditioning tensor y (first-frame conditioning mask)
    img_lat = initial_latent.permute(0, 2, 1, 3, 4)
    total_frames = num_frames + 20
    msk = torch.zeros_like(img_lat.repeat(1, 1, total_frames, 1, 1)[:, :1])
    image_cat = img_lat.repeat(1, 1, total_frames, 1, 1)
    msk[:, :, 1:] = 1
    y = torch.cat([image_cat, msk], dim=1)

    # Audio embeddings
    audio_emb = encode_audio(audio_path, num_frames=num_frames).to(device=device, dtype=torch.bfloat16)

    # Noise tensor
    h, w = initial_latent.shape[-2], initial_latent.shape[-1]
    noise = torch.randn((1, num_frames - 1, 16, h, w), device=device, dtype=torch.bfloat16)

    with torch.no_grad():
        video = pipeline.inference_avatar_forcing(
            noise=noise,
            text_prompts=[prompt],
            audio_embeddings=audio_emb,
            y=y,
            return_latents=False,
            initial_latent=initial_latent,
        )

    pipeline.vae.model.clear_cache()

    # Decode: [B, T, C, H, W] → list of [H, W, C] uint8
    frames_np = (255.0 * rearrange(video[0], "t c h w -> t h w c")).cpu().numpy().astype("uint8")
    raw_path = tempfile.mktemp(suffix=".mp4")
    writer = imageio.get_writer(raw_path, fps=25, codec="libx264", quality=8)
    for frame in frames_np:
        writer.append_data(frame)
    writer.close()

    # Mux original audio track
    out_path = tempfile.mktemp(suffix=".mp4")
    subprocess.run([
        "ffmpeg", "-y",
        "-i", raw_path,
        "-i", audio_path,
        "-c:v", "copy", "-c:a", "aac", "-shortest",
        out_path,
    ], check=True, capture_output=True)
    os.unlink(raw_path)

    return out_path


# ── Gradio UI ─────────────────────────────────────────────────────────────────
with gr.Blocks(title="AvatarForcing · AIBRUH") as demo:
    gr.Markdown("## AvatarForcing · Streaming Talking Avatar\n*Portrait + Speech → Animated MP4 @ 25 FPS*")
    with gr.Row():
        portrait_in = gr.Image(type="filepath", label="Portrait (JPG/PNG)")
        audio_in    = gr.Audio(type="filepath", label="Speech Audio (WAV/MP3)")
    prompt_in  = gr.Textbox(
        value="A photorealistic person speaking naturally, warm cinematic lighting, shallow depth of field, ultra detailed",
        label="Text Prompt",
    )
    seconds_in = gr.Slider(1, 10, value=5, step=1, label="Duration (seconds)")
    btn        = gr.Button("Generate Avatar Video", variant="primary")
    video_out  = gr.Video(label="Amanda Speaking")

    btn.click(
        fn=generate,
        inputs=[portrait_in, audio_in, prompt_in, seconds_in],
        outputs=video_out,
        api_name="generate",
    )

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