# 启用bash的扩展功能
shopt -s extglob

# 检查参数是否传入
if [ $# -eq 0 ]; then
    echo "Error: Insufficient parameters"
    echo
    echo "Usage:"
    echo "  ./run_train_and_cloneb.sh \"input_video\""
    echo
    echo "Examples:"
    echo "  ./run_train_and_cloneb.sh \"path/to/your/video.mp4\""
    exit 1
fi

# 定义核心变量
INPUT_VIDEO="$1"

# 获取脚本所在目录
SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/
cd "$SCRIPT_DIR" || exit 1

# 使用固定的预生成 avatar 目录（不生成时间戳）
# 如果需要生成新的，取消下面的注释并注释掉固定名称
AVATAR_ID="wav2lip256_avatar1"
# TIMESTAMP=$(date +%Y%m%d_%H%M%S)
# AVATAR_ID="avatar_${TIMESTAMP}"
AVATAR_DIR="${SCRIPT_DIR}data/avatars/${AVATAR_ID}"

# 不自动删除 avatar 目录（因为是预生成的）
# cleanup_avatar_dir() {
#     if [ -n "${AVATAR_DIR:-}" ] && [ -d "${AVATAR_DIR}" ]; then
#         echo "Cleaning generated avatar directory: ${AVATAR_DIR}"
#         rm -rf -- "${AVATAR_DIR}"
#     fi
# }
# trap 'cleanup_avatar_dir' EXIT

# 默认环境变量（可通过外部传入覆盖）
: "${CUDA_VISIBLE_DEVICES:=4,5,7}"
: "${AIOICE_PORT_MIN:=50000}"
: "${AIOICE_PORT_MAX:=50010}"
: "${AIOICE_BIND_IP:=192.168.22.9}"
: "${WEBRTC_NAT_IP:=183.252.196.135}"
# PyTorch 显存优化
: "${PYTORCH_CUDA_ALLOC_CONF:=max_split_size_mb:512,expandable_segments:True}"
# 禁用 TorchDynamo 编译（避免 VoxCPM2 兼容性问题）
: "${TORCHDYNAMO_DISABLE:=1}"
# RTMP 推流配置（使用 nginx-rtmp + FFmpeg）
: "${ENABLE_RTMP:=true}"  # 是否启用 RTMP 推流（默认启用）
: "${RTMP_PORT:=1935}"    # RTMP 服务端口
: "${FFMPEG_PUSH:=true}"  # 是否启动本机 FFmpeg 转发
export CUDA_VISIBLE_DEVICES AIOICE_PORT_MIN AIOICE_PORT_MAX AIOICE_BIND_IP WEBRTC_NAT_IP
export PYTORCH_CUDA_ALLOC_CONF TORCHDYNAMO_DISABLE

echo "========================================"
echo "Digital Human Generation Pipeline (Green Screen)"
echo "========================================"
echo
echo "Input Video: ${INPUT_VIDEO}"
echo "Avatar ID: ${AVATAR_ID}"
echo "Timestamp: ${TIMESTAMP}"
echo "Using env: CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES} AIOICE_BIND_IP=${AIOICE_BIND_IP} WEBRTC_NAT_IP=${WEBRTC_NAT_IP} AIOICE_PORT_MIN=${AIOICE_PORT_MIN} AIOICE_PORT_MAX=${AIOICE_PORT_MAX}"
echo

# 检查输入视频是否存在
if [ ! -f "${INPUT_VIDEO}" ]; then
    echo "Error: Input video not found: ${INPUT_VIDEO}"
    exit 1
fi

echo "Current Directory: $(pwd)"
echo

# ========================================
# Step 1: Convert video background to green screen
# ========================================
echo
echo "[Step 1/3] Converting video background to green screen..."
echo

# 确保 back/video 目录存在
mkdir -p "${SCRIPT_DIR}back/video"

# 生成输出文件名
INPUT_BASENAME=$(basename "${INPUT_VIDEO}" | sed 's/\.[^.]*$//')
GREENSCREEN_VIDEO="${SCRIPT_DIR}back/video/${INPUT_BASENAME}_greenscreen.mp4"

echo "Converting video to green screen..."
echo "  Input: ${INPUT_VIDEO}"
echo "  Output: ${GREENSCREEN_VIDEO}"
echo

# 运行绿幕转换脚本
python back/video_to_greenscreen.py "${INPUT_VIDEO}" -o "${GREENSCREEN_VIDEO}"

if [ $? -ne 0 ]; then
    echo "Error: Green screen conversion failed!"
    cd "${SCRIPT_DIR}" || exit 1
    exit 1
fi

# 检查输出文件是否存在
if [ ! -f "${GREENSCREEN_VIDEO}" ]; then
    echo "Error: Green screen video not generated!"
    cd "${SCRIPT_DIR}" || exit 1
    exit 1
fi

echo
echo "Green screen conversion completed!"
echo "Generated Video: ${GREENSCREEN_VIDEO}"
echo

# ========================================
# Step 2: Run genavatar.py for pre-processing
# ========================================
echo
echo "[Step 2/3] Running genavatar.py for pre-processing..."
echo

# 进入wav2lip目录
cd "${SCRIPT_DIR}wav2lip" || exit 1

echo "Generating avatar data with:"
echo "  Video: ${GREENSCREEN_VIDEO}"
echo "  Avatar ID: ${AVATAR_ID}"
echo "  Image Size: 256"
echo

# 运行 genavatar.py
python genavatar.py --video_path "${GREENSCREEN_VIDEO}" --img_size 256 --avatar_id "${AVATAR_ID}"

if [ $? -ne 0 ]; then
    echo "Error: genavatar.py failed!"
    cd "${SCRIPT_DIR}" || exit 1
    exit 1
fi

cd "${SCRIPT_DIR}" || exit 1

echo
echo "Pre-processing completed!"
echo "Avatar data saved to: data/avatars/${AVATAR_ID}"
echo

# ========================================
# Step 3: 启动 VoxCPM2 API 服务 + 运行 app.py
# ========================================
echo
echo "[Step 3/4] 检查并启动 VoxCPM2 API 服务..."
echo

# VoxCPM2 API 服务配置
VOXCPM2_API_URL="http://localhost:6003"
VOXCPM2_API_LOG="/tmp/voxcpm2_api.log"

echo "检查 VoxCPM2 API 服务状态..."
echo "  API 地址：${VOXCPM2_API_URL}"

# 检查 API 服务是否已在运行
if curl -s ${VOXCPM2_API_URL}/health > /dev/null 2>&1; then
    echo "✅ VoxCPM2 API 服务已在运行"
    curl -s ${VOXCPM2_API_URL}/health | python -m json.tool 2>/dev/null || true
else
    echo "🚀 VoxCPM2 API 服务未运行，正在启动..."
    echo
    
    # 检查模型目录是否存在
    if [ ! -d "/mnt/nvme1data/model/VoxCPM2" ]; then
        echo "❌ 错误：VoxCPM2 模型目录不存在：/mnt/nvme1data/model/VoxCPM2"
        echo "请先下载或复制模型文件到该目录"
        exit 1
    fi
    
    # 检查 API 脚本是否存在
    if [ ! -f "/mnt/nvme1data/model/voxcpm2_api.py" ]; then
        echo "❌ 错误：VoxCPM2 API 脚本不存在：/mnt/nvme1data/model/voxcpm2_api.py"
        exit 1
    fi
    
    # 在后台启动 API 服务
    cd /mnt/nvme1data/model
    nohup python voxcpm2_api.py > ${VOXCPM2_API_LOG} 2>&1 &
    API_PID=$!
    echo "✅ VoxCPM2 API 服务已启动"
    echo "   PID: ${API_PID}"
    echo "   日志：${VOXCPM2_API_LOG}"
    echo
    
    # 等待服务启动（最多 120 秒）
    echo "等待模型加载（最多120秒）..."
    START_TIME=$(date +%s)
    
    for i in $(seq 1 120); do
        if curl -s ${VOXCPM2_API_URL}/health > /dev/null 2>&1; then
            END_TIME=$(date +%s)
            WAIT_TIME=$((END_TIME - START_TIME))
            echo "✅ VoxCPM2 API 服务已就绪（等待 ${WAIT_TIME} 秒）"
            echo
            # 显示服务信息
            curl -s ${VOXCPM2_API_URL}/health | python -m json.tool 2>/dev/null || true
            break
        fi
        
        # 每 10 秒显示一次进度
        if [ $((i % 10)) -eq 0 ]; then
            echo "   已等待 ${i} 秒..."
        fi
        
        if [ $i -eq 120 ]; then
            echo
            echo "❌ VoxCPM2 API 服务启动超时（120秒）"
            echo
            echo "最后 50 行日志："
            echo "========================================"
            tail -50 ${VOXCPM2_API_LOG}
            echo "========================================"
            echo
            echo "请检查："
            echo "  1. 模型文件是否完整：ls -lh /mnt/nvme1data/model/VoxCPM2/"
            echo "  2. voxcpm 包是否安装：python -c 'import voxcpm'"
            echo "  3. GPU 是否可用：nvidia-smi"
            echo
            exit 1
        fi
        sleep 1
    done
    
    cd "${SCRIPT_DIR}"
fi

echo
echo "========================================"
echo "[Step 4/4] 启动数字人服务（API 调用模式）"
echo "========================================"
echo

echo "Starting digital human server..."
echo "Avatar ID: ${AVATAR_ID}"
echo "TTS: voxcpm2api (VoxCPM2 API 调用模式 - 显存隔离)"
echo "API URL: ${VOXCPM2_API_URL}"
echo "Model: /mnt/nvme1data/model/VoxCPM2 (独立进程)"
echo "Transport: ${TRANSPORT_MODE}"
if [ "${ENABLE_RTMP}" = "true" ]; then
    echo "RTMP Push URL: ${PUSH_URL}"
    echo "Jump Server Command:"
    echo "  ffmpeg -i rtmp://127.0.0.1:1935/live/${AVATAR_ID} -c copy -f flv \"rtmp://livepush.ailien.shop/alien/people?auth_key=...\""
fi
echo
echo "优势："
echo "  ✅ 模型显存独立（~8GB），不影响数字人渲染"
echo "  ✅ 多个数字人实例可共用同一 API 服务"
echo "  ✅ API 崩溃不影响数字人主进程"
echo "  ✅ 更新模型只需重启 API，无需停数字人"
echo
echo "访问地址："
echo "  Web 界面：http://localhost:7868/dashboard.html"
echo "  API 健康检查：${VOXCPM2_API_URL}/health"
echo
echo "Press Ctrl+C to stop the server."
echo

# 检查并安装pyaudio
echo "Checking pyaudio installation..."
python -c "import pyaudio" 2>/dev/null || (echo "Installing pyaudio..." && pip install pyaudio)

# ========================================
# Step 3.1: 检查并配置 nginx-rtmp（如果需要）
# ========================================
if [ "${ENABLE_RTMP}" = "true" ]; then
    echo
    echo "[Step 3.1/4] Checking nginx-rtmp configuration..."
    echo
    
    # 检查 nginx RTMP 模块是否已安装
    if dpkg -l | grep -q libnginx-mod-rtmp; then
        echo "✅ nginx-rtmp-module is already installed"
    else
        echo "⚠️  nginx-rtmp-module not found. Please run:"
        echo "   sudo bash setup_nginx_rtmp.sh"
        echo
        echo "Continuing without RTMP support..."
        ENABLE_RTMP="false"
    fi
    
    if [ "${ENABLE_RTMP}" = "true" ]; then
        # 检查 nginx 是否运行
        if systemctl is-active --quiet nginx; then
            echo "✅ nginx is running"
        else
            echo "Starting nginx..."
            sudo systemctl start nginx
            if [ $? -eq 0 ]; then
                echo "✅ nginx started successfully"
            else
                echo "❌ Failed to start nginx"
                ENABLE_RTMP="false"
            fi
        fi
        
        # 设置推流地址
        # app.py 推送到 nginx-rtmp，跳板机从 nginx-rtmp 拉流
        PUSH_URL="rtmp://127.0.0.1:${RTMP_PORT}/live/${AVATAR_ID}"
        TRANSPORT_MODE="rtcpush"
        
        echo
        echo "RTMP Configuration:"
        echo "  Push URL: ${PUSH_URL}"
        echo "  Jump Server Pull: ffmpeg -i rtmp://127.0.0.1:1935/live/${AVATAR_ID} -c copy -f flv \"rtmp://cdn...\""
        echo
    fi
else
    PUSH_URL=""
    TRANSPORT_MODE="webrtc"
fi

# 启动 app.py
# 关键环境变量说明：
#   AIOICE_PORT_MIN/MAX  : aiortc 绑定的 UDP 端口范围，必须与防火墙放行范围一致
#   AIOICE_BIND_IP       : aiortc 只绑定这一个 IP（服务器有多网卡时，指定浏览器能路由到的那个）
#   WEBRTC_NAT_IP        : 如果浏览器用的是跳板机 IP 而非服务器 IP，设置此项做 SDP IP 替换
#
# 默认值（可通过外部传入环境变量覆盖）：
#   CUDA_VISIBLE_DEVICES=6,7  (使用 GPU 6,7 运行)
#   AIOICE_PORT_MIN=50000  AIOICE_PORT_MAX=50010
#   AIOICE_BIND_IP=192.168.22.9  WEBRTC_NAT_IP=183.252.196.135
#
echo "Using env: CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES} AIOICE_BIND_IP=${AIOICE_BIND_IP} WEBRTC_NAT_IP=${WEBRTC_NAT_IP} AIOICE_PORT_MIN=${AIOICE_PORT_MIN} AIOICE_PORT_MAX=${AIOICE_PORT_MAX}"

CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES}" \
AIOICE_BIND_IP="${AIOICE_BIND_IP}" \
WEBRTC_NAT_IP="${WEBRTC_NAT_IP}" \
AIOICE_PORT_MIN="${AIOICE_PORT_MIN}" \
AIOICE_PORT_MAX="${AIOICE_PORT_MAX}" \
python app.py \
    --avatar_id "${AVATAR_ID}" \
    --tts voxcpm2api \
    --VOXCPM2_API_URL ${VOXCPM2_API_URL} \
    --VOXCPM2_REF_WAV voice_output.wav \
    --VOXCPM2_REF_TEXT "你好，买水果，卖水果，新鲜的水果。" \
    --CFG_VALUE 1.5 \
    --INFERENCE_TIMESTEPS 10 \
    --model wav2lip \
    --transport ${TRANSPORT_MODE} \
    ${PUSH_URL:+--push_url "${PUSH_URL}"} \
    --listenport 7868

echo
echo "========================================"
echo "Pipeline completed!"
echo "========================================"