#!/bin/bash
#
# Download MuseTalk models from HuggingFace
#

set -e

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
PROJECT_DIR="$(dirname "$SCRIPT_DIR")"
MODELS_DIR="${PROJECT_DIR}/models"

GREEN='\033[0;32m'
CYAN='\033[0;36m'
YELLOW='\033[1;33m'
NC='\033[0m'

success() { echo -e "${GREEN}✓${NC} $1"; }
header() { echo -e "\n${CYAN}$1${NC}"; }
warn() { echo -e "${YELLOW}!${NC} $1"; }

echo ""
echo "╔══════════════════════════════════════════════════════════════╗"
echo "║          MUSETALK - MODEL DOWNLOAD                           ║"
echo "╚══════════════════════════════════════════════════════════════╝"
echo ""

# Check for huggingface-cli
if ! command -v huggingface-cli &> /dev/null; then
    echo "Installing huggingface_hub..."
    pip install huggingface_hub -q
fi

# Create directories
mkdir -p "${MODELS_DIR}/avatar/musetalk"
mkdir -p "${MODELS_DIR}/avatar/vae"
mkdir -p "${MODELS_DIR}/whisper"
mkdir -p "${MODELS_DIR}/dwpose"
mkdir -p "${MODELS_DIR}/face_parse"

# ============================================================
# 1. MUSETALK MODELS
# ============================================================
header "[1/5] Downloading MuseTalk UNet..."

if [ -f "${MODELS_DIR}/avatar/musetalk/pytorch_model.bin" ]; then
    success "MuseTalk UNet already exists"
else
    huggingface-cli download TMElyralab/MuseTalk \
        --include "models/musetalk/*" \
        --local-dir "${MODELS_DIR}/temp_musetalk" \
        --local-dir-use-symlinks False

    # Move files to correct location
    if [ -d "${MODELS_DIR}/temp_musetalk/models/musetalk" ]; then
        cp -r "${MODELS_DIR}/temp_musetalk/models/musetalk/"* "${MODELS_DIR}/avatar/musetalk/"
        rm -rf "${MODELS_DIR}/temp_musetalk"
        success "MuseTalk UNet downloaded"
    else
        warn "MuseTalk download may have failed"
    fi
fi

# ============================================================
# 2. VAE (SD-VAE-FT-MSE)
# ============================================================
header "[2/5] Downloading VAE..."

if [ -f "${MODELS_DIR}/avatar/vae/config.json" ]; then
    success "VAE already exists"
else
    huggingface-cli download stabilityai/sd-vae-ft-mse \
        --local-dir "${MODELS_DIR}/avatar/vae" \
        --local-dir-use-symlinks False

    success "VAE downloaded"
fi

# ============================================================
# 3. WHISPER
# ============================================================
header "[3/5] Downloading Whisper..."

if [ -f "${MODELS_DIR}/whisper/tiny.pt" ]; then
    success "Whisper already exists"
else
    # Download from OpenAI
    wget -q -O "${MODELS_DIR}/whisper/tiny.pt" \
        "https://openaipublic.azureedge.net/main/whisper/models/65147644a518d12f04e32d6f3b26facc3f8dd46e5390956a9424a650c0ce22b9/tiny.pt" \
        || curl -s -o "${MODELS_DIR}/whisper/tiny.pt" \
        "https://openaipublic.azureedge.net/main/whisper/models/65147644a518d12f04e32d6f3b26facc3f8dd46e5390956a9424a650c0ce22b9/tiny.pt"

    if [ -f "${MODELS_DIR}/whisper/tiny.pt" ]; then
        success "Whisper downloaded"
    else
        warn "Whisper download failed"
    fi
fi

# ============================================================
# 4. DWPOSE
# ============================================================
header "[4/5] Downloading DWPose..."

if [ -f "${MODELS_DIR}/dwpose/dw-ll_ucoco_384.pth" ]; then
    success "DWPose already exists"
else
    huggingface-cli download yzd-v/DWPose \
        --include "dw-ll_ucoco_384.pth" \
        --local-dir "${MODELS_DIR}/dwpose" \
        --local-dir-use-symlinks False

    success "DWPose downloaded"
fi

# ============================================================
# 5. FACE PARSING
# ============================================================
header "[5/5] Downloading Face Parsing models..."

if [ -f "${MODELS_DIR}/face_parse/79999_iter.pth" ]; then
    success "Face parsing already exists"
else
    # Download from MuseTalk repo
    huggingface-cli download TMElyralab/MuseTalk \
        --include "models/face-parse-bisenet/*" \
        --local-dir "${MODELS_DIR}/temp_fp" \
        --local-dir-use-symlinks False

    if [ -d "${MODELS_DIR}/temp_fp/models/face-parse-bisenet" ]; then
        cp -r "${MODELS_DIR}/temp_fp/models/face-parse-bisenet/"* "${MODELS_DIR}/face_parse/"
        rm -rf "${MODELS_DIR}/temp_fp"
        success "Face parsing downloaded"
    fi

    # Also need resnet18
    if [ ! -f "${MODELS_DIR}/face_parse/resnet18-5c106cde.pth" ]; then
        wget -q -O "${MODELS_DIR}/face_parse/resnet18-5c106cde.pth" \
            "https://download.pytorch.org/models/resnet18-5c106cde.pth" \
            || curl -s -o "${MODELS_DIR}/face_parse/resnet18-5c106cde.pth" \
            "https://download.pytorch.org/models/resnet18-5c106cde.pth"
    fi
fi

# ============================================================
# SUMMARY
# ============================================================
echo ""
echo "╔══════════════════════════════════════════════════════════════╗"
echo "║          DOWNLOAD COMPLETE                                   ║"
echo "╚══════════════════════════════════════════════════════════════╝"
echo ""

echo "Models directory: ${MODELS_DIR}"
echo ""

# Check all models
check_model() {
    if [ -f "$1" ] || [ -d "$1" ]; then
        echo -e "  ${GREEN}✓${NC} $2"
    else
        echo -e "  ✗ $2 (missing)"
    fi
}

echo "Installed models:"
check_model "${MODELS_DIR}/avatar/musetalk/pytorch_model.bin" "MuseTalk UNet"
check_model "${MODELS_DIR}/avatar/vae/config.json" "VAE"
check_model "${MODELS_DIR}/whisper/tiny.pt" "Whisper Tiny"
check_model "${MODELS_DIR}/dwpose/dw-ll_ucoco_384.pth" "DWPose"
check_model "${MODELS_DIR}/face_parse/79999_iter.pth" "Face Parsing"
echo ""