#!/usr/bin/env bash
set -e

# ==============================================================================
# SadTalker Ultimate Setup & Execution Cheatsheet
# Official Repo: https://github.com/OpenTalker/SadTalker
# Reference Machine Tested: NVIDIA GeForce RTX 4090 24GB, Ubuntu 22.04 / Vast.ai
# ==============================================================================

echo "=== [1/7] Checking & Installing System Prerequisites ==="
if command -v apt-get &> /dev/null; then
    echo "Ubuntu/Debian detected. Updating apt packages..."
    apt update && apt install -y git ffmpeg wget python3-pip
elif command -v brew &> /dev/null; then
    echo "macOS detected. Installing ffmpeg via brew if missing..."
    brew install ffmpeg || true
fi

echo "=== [2/7] Cloning SadTalker Repository ==="
if [ ! -d "SadTalker" ]; then
    git clone https://github.com/OpenTalker/SadTalker.git
fi
cd SadTalker

echo "=== [3/7] Setting Up Python Environment ==="
# Verify or create conda env if conda is installed
if command -v conda &> /dev/null; then
    echo "Conda detected. Ensuring environment 'sadtalker' with Python 3.8 exists..."
    conda create -n sadtalker python=3.8 -y || true
    echo "Activate it with: conda activate sadtalker"
fi

echo "=== [4/7] Installing PyTorch & Project Dependencies ==="
pip install torch torchvision torchaudio
pip install -r requirements.txt

# Fallback prerequisite: FFmpeg inside conda if missing from PATH
if ! command -v ffmpeg &> /dev/null && command -v conda &> /dev/null; then
    echo "FFmpeg binary not detected in PATH. Installing via conda-forge..."
    conda install -c conda-forge ffmpeg -y
fi

echo "=== [5/7] Applying Undocumented Bug Fixes ==="
# Fix 1: ModuleNotFoundError: No module named 'torchvision.transforms.functional_tensor'
python3 -c "
import os, sys
site_packages = [p for p in sys.path if 'site-packages' in p]
for sp in site_packages:
    fpath = os.path.join(sp, 'basicsr', 'data', 'degradations.py')
    if os.path.exists(fpath):
        with open(fpath, 'r') as f:
            c = f.read()
        if 'torchvision.transforms.functional_tensor' in c:
            c = c.replace('from torchvision.transforms.functional_tensor import rgb_to_grayscale', 'from torchvision.transforms.functional import rgb_to_grayscale')
            with open(fpath, 'w') as f:
                f.write(c)
            print('[SUCCESS] Patched basicsr degradations.py for torchvision compatibility.')
"

# Fix 2: Gradio version compatibility (avoids 'Row object has no attribute style')
pip install "gradio==3.41.2" TTS==0.13.3

echo "=== [6/7] Downloading Model Checkpoints ==="
bash scripts/download_models.sh

echo "=== [7/7] Installation Complete! ==="
echo ""
echo "To launch Gradio WebUI: python app_sadtalker.py"
echo ""
echo "To run CLI inference:"
echo "python inference.py \\"
echo "  --driven_audio ./examples/driven_audio/bus_chinese.wav \\"
echo "  --source_image ./examples/source_image/full_body_1.png \\"
echo "  --result_dir ./results \\"
echo "  --enhancer gfpgan"