import os
from typing import Dict, Any
from PIL import Image
import tempfile
import base64
import shutil

from src.gradio_demo import SadTalker
from huggingface_hub import snapshot_download


# Download model checkpoints at init
REPO_ID = "vinthony/SadTalker-V002rc"
snapshot_download(
    repo_id=REPO_ID, local_dir="./checkpoints", local_dir_use_symlinks=True
)

# Load SadTalker
sad_talker = SadTalker(lazy_load=True)


def save_base64_to_file(base64_str, suffix):
    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
    data = base64.b64decode(base64_str)
    with open(temp_file.name, "wb") as f:
        f.write(data)
    return temp_file.name


def predict(inputs: Dict[str, Any]) -> Dict[str, Any]:
    """
    Expected input JSON:
    {
        "image": base64_image_string,
        "audio": base64_audio_string,
        "ref_video": base64_video_string or None,
        "use_ref_video": true/false,
        "use_idle_mode": true/false,
        "length": 5,
        "pose_style": 0,
        "expression_scale": 1.0,
        "preprocess": "crop",
        ...
    }
    """
    try:
        image_path = save_base64_to_file(inputs["image"], ".png")
        audio_path = (
            save_base64_to_file(inputs["audio"], ".wav")
            if not inputs.get("use_idle_mode", False)
            else None
        )
        ref_video_path = (
            save_base64_to_file(inputs["ref_video"], ".mp4")
            if inputs.get("use_ref_video")
            else None
        )

        output_path = sad_talker.test(
            source_image=image_path,
            driven_audio=audio_path,
            preprocess=inputs.get("preprocess", "crop"),
            still_mode=inputs.get("still_mode", False),
            enhancer=inputs.get("enhancer", False),
            batch_size=inputs.get("batch_size", 1),
            size_of_image=inputs.get("size", 256),
            pose_style=inputs.get("pose_style", 0),
            facerender=inputs.get("facerender", "facevid2vid"),
            expression_scale=inputs.get("expression_scale", 1.0),
            use_ref_video=inputs.get("use_ref_video", False),
            ref_video=ref_video_path,
            ref_info=inputs.get("ref_info", "pose"),
            use_idle_mode=inputs.get("use_idle_mode", False),
            length_of_audio=inputs.get("length", 5),
            blink_every=inputs.get("blink", True),
        )

        with open(output_path, "rb") as f:
            video_bytes = f.read()

        video_b64 = base64.b64encode(video_bytes).decode("utf-8")

        return {"video": video_b64}

    except Exception as e:
        return {"error": str(e)}