"""
AIBRUH/ditto — Amanda's talking head engine.
SadTalker ZeroGPU A100: portrait image + audio → MP4 talking head video.
"""
import gradio as gr
import spaces
import os, sys, subprocess

SADTALKER = "/tmp/SadTalker"
CKPT = f"{SADTALKER}/checkpoints"

def _setup():
    # Clone SadTalker source if not already present
    if not os.path.exists(SADTALKER):
        subprocess.run(
            ["git", "clone", "--depth=1",
             "https://github.com/OpenTalker/SadTalker.git", SADTALKER],
            check=True,
        )

    # Install ALL SadTalker dependencies in one shot from their own requirements file
    req_file = os.path.join(SADTALKER, "requirements.txt")
    if os.path.exists(req_file):
        subprocess.run(
            [sys.executable, "-m", "pip", "install", "-q", "-r", req_file],
            check=False,
        )

    # numpy 2.x removed VisibleDeprecationWarning — patch it now (before any SadTalker import)
    import numpy as np
    if not hasattr(np, "VisibleDeprecationWarning"):
        np.VisibleDeprecationWarning = DeprecationWarning  # type: ignore[attr-defined]

    # Download SadTalker checkpoints if not present
    if not os.path.exists(f"{CKPT}/SadTalker_V0.0.2_256.safetensors"):
        os.makedirs(CKPT, exist_ok=True)
        from huggingface_hub import snapshot_download
        snapshot_download(
            repo_id="vinthony/SadTalker-V002rc",
            local_dir=CKPT,
            ignore_patterns=["*.md", "*.txt"],
        )

    if SADTALKER not in sys.path:
        sys.path.insert(0, SADTALKER)


_setup()

# numpy patch must also run at module level for the import below
import numpy as np  # noqa: E402
if not hasattr(np, "VisibleDeprecationWarning"):
    np.VisibleDeprecationWarning = DeprecationWarning  # type: ignore[attr-defined]

from src.gradio_demo import SadTalker as _ST  # noqa: E402

_sad: _ST | None = None

def _load():
    global _sad
    if _sad is None:
        _sad = _ST(CKPT, f"{SADTALKER}/config", lazy_load=True)
    return _sad


@spaces.GPU
def infer(source_image: str, driven_audio: str):
    """Portrait image + WAV audio → talking head MP4."""
    st = _load()
    result = st.test(
        source_image=source_image,
        driven_audio=driven_audio,
        preprocess="crop",
        still_mode=True,
        use_enhancer=False,
        batch_size=1,
        size=256,
        pose_style=0,
        exp_scale=1.0,
        use_ref_video=False,
        ref_video=None,
        ref_info="pose",
        use_idle_mode=False,
        length_of_audio=0,
        use_blink=True,
    )
    return result


with gr.Blocks(title="AIBRUH/ditto · Amanda Engine") as demo:
    gr.Markdown("## AIBRUH/ditto · Amanda Talking Head Engine")
    with gr.Row():
        img_in = gr.Image(type="filepath", label="Source Portrait (JPG/PNG)")
        aud_in = gr.Audio(type="filepath", label="Driving Audio (WAV)")
    btn = gr.Button("Generate Talking Head", variant="primary")
    vid_out = gr.Video(label="Amanda Speaking")
    btn.click(fn=infer, inputs=[img_in, aud_in], outputs=vid_out, api_name="infer")

demo.launch()