"""Nova Sonic Avatar · ZeroGPU Space
Sonic (CVPR 2025) audio→portrait talking-head video.
Fork of github.com/jixiaozhong/Sonic adapted for HF ZeroGPU lazy-load.
"""
from __future__ import annotations

import hashlib
import os
import subprocess
import time
from pathlib import Path

import gradio as gr
import numpy as np
import torch
from pydub import AudioSegment

try:
    import spaces
except Exception:
    class _Spaces:
        def GPU(self, *args, **kwargs):
            def deco(fn):
                return fn
            return deco
    spaces = _Spaces()

# Paths the Sonic pipeline expects
ROOT = Path(__file__).parent.resolve()
CKPT_DIR = ROOT / "checkpoints"
TMP_DIR = ROOT / "tmp_path"
RES_DIR = ROOT / "res_path"
TMP_DIR.mkdir(exist_ok=True)
RES_DIR.mkdir(exist_ok=True)


def _ensure_checkpoints():
    """Pull Sonic + SVD-XT + whisper-tiny from HF on first call."""
    if (CKPT_DIR / "Sonic" / "unet.pth").exists():
        return
    print("[sonic] downloading checkpoints (first call · ~12GB)...", flush=True)
    from huggingface_hub import snapshot_download
    CKPT_DIR.mkdir(exist_ok=True)
    snapshot_download("LeonJoe13/Sonic", local_dir=str(CKPT_DIR))
    snapshot_download("stabilityai/stable-video-diffusion-img2vid-xt", local_dir=str(CKPT_DIR / "stable-video-diffusion-img2vid-xt"))
    snapshot_download("openai/whisper-tiny", local_dir=str(CKPT_DIR / "whisper-tiny"))
    print("[sonic] checkpoints ready", flush=True)


_PIPE = None


def _load_pipe():
    global _PIPE
    if _PIPE is not None:
        return _PIPE
    _ensure_checkpoints()
    from sonic import Sonic
    print("[sonic] loading pipeline on cuda...", flush=True)
    _PIPE = Sonic(0)
    print("[sonic] pipeline ready", flush=True)
    return _PIPE


def _md5(content: bytes) -> str:
    return hashlib.md5(content).hexdigest()


@spaces.GPU(duration=300)
def render_avatar(image, audio, dynamic_scale: float = 1.0) -> tuple[str | None, str]:
    """
    image: PIL or filepath
    audio: gradio Audio tuple (sample_rate, np.array)
    dynamic_scale: motion intensity
    returns: (mp4_path, status_json)
    """
    import json as _json
    started = time.time()
    pipe = _load_pipe()

    # Persist image to disk (Sonic API takes paths)
    if isinstance(image, str):
        img_arr = np.array([])
        img_path = image
    else:
        img_arr = np.array(image)
        img_md5 = _md5(img_arr.tobytes())
        img_path = str(TMP_DIR / f"{img_md5}.png")
        from PIL import Image as _Image
        _Image.fromarray(img_arr).save(img_path)

    # Persist audio to disk as wav
    sampling_rate, arr = audio[:2]
    if len(arr.shape) == 1:
        arr = arr[:, None]
    seg = AudioSegment(
        arr.tobytes(),
        frame_rate=sampling_rate,
        sample_width=arr.dtype.itemsize,
        channels=arr.shape[1],
    )
    audio_md5 = _md5(seg.raw_data)
    audio_path = str(TMP_DIR / f"{audio_md5}.wav")
    seg.export(audio_path, format="wav")

    res_video_path = str(RES_DIR / f"{audio_md5}_{dynamic_scale}.mp4")
    if os.path.isfile(res_video_path):
        meta = {"cached": True, "wall_ms": int((time.time() - started) * 1000)}
        return res_video_path, _json.dumps(meta, indent=2)

    expand_ratio = 0.5
    min_resolution = 512
    inference_steps = 25

    face_info = pipe.preprocess(img_path, expand_ratio=expand_ratio)
    if face_info.get("face_num", 0) <= 0:
        meta = {"error": "no face detected", "face_info": face_info, "wall_ms": int((time.time() - started) * 1000)}
        return None, _json.dumps(meta, indent=2)

    crop_path = img_path + ".crop.png"
    pipe.crop_image(img_path, crop_path, face_info["crop_bbox"])
    pipe.process(
        crop_path,
        audio_path,
        res_video_path,
        min_resolution=min_resolution,
        inference_steps=inference_steps,
        dynamic_scale=dynamic_scale,
    )
    meta = {
        "cached": False,
        "inference_steps": inference_steps,
        "min_resolution": min_resolution,
        "dynamic_scale": dynamic_scale,
        "wall_ms": int((time.time() - started) * 1000),
    }
    return res_video_path, _json.dumps(meta, indent=2)


with gr.Blocks(title="Nova Sonic Avatar") as demo:
    gr.Markdown("# Nova Sonic Avatar · ZeroGPU")
    gr.Markdown("Audio → talking-head video (Sonic, CVPR 2025). Backs `Talk to Nova` cockpit per-turn render.")
    with gr.Row():
        img_in = gr.Image(label="Reference portrait", type="pil")
        aud_in = gr.Audio(label="Speech audio", type="numpy")
    scale = gr.Slider(0.5, 2.0, value=1.0, step=0.1, label="dynamic_scale (motion intensity)")
    vid_out = gr.Video(label="Talking-head video")
    meta_out = gr.Code(label="Render metadata", language="json")
    gr.Button("Render").click(render_avatar, [img_in, aud_in, scale], [vid_out, meta_out])

if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1).launch(server_name="0.0.0.0", server_port=7860)