#!/usr/bin/env bash
# =============================================================================
# Trading Studio 一键部署脚本
# 安装路径: /opt/Trading_Studio
# 运行用户: root
# 功能: 系统依赖 → 代码拉取 → Python 虚拟环境 → PyTorch CUDA → PM2 常驻
#
# 用法:
#   sudo bash deploy.sh              # 首次完整部署并 PM2 启动
#   sudo bash deploy.sh update       # 拉取最新代码、更新依赖、重启 PM2
#   sudo bash deploy.sh restart      # 仅重启 PM2 进程
#   sudo bash deploy.sh stop         # 停止 PM2 进程
#   sudo bash deploy.sh status       # 查看 PM2 与 GPU 状态
#   sudo bash deploy.sh logs         # 查看 PM2 最近日志
#
# 环境变量:
#   USE_CN_MIRROR=1   使用清华 PyPI / PyTorch 镜像（国内服务器推荐，默认开启）
#   USE_CN_MIRROR=0   使用官方 PyTorch 源
#   PIP_TIMEOUT=600   pip 单次下载超时秒数（默认 600）
#   SKIP_PYTORCH=1    跳过 PyTorch 安装（已手动装好时使用）
# =============================================================================

set -euo pipefail

# ---------------------------------------------------------------------------
# 可配置常量
# ---------------------------------------------------------------------------
INSTALL_DIR="/opt/Trading_Studio"
GIT_REPO="https://git.bz121.com/dekun/Trading_Studio.git"
GIT_BRANCH="main"
PM2_APP_NAME="trading_studio"
GRADIO_PORT=5683
GPU_POWER_LIMIT=120

# pip 网络参数（大包如下载 triton 需要更长超时）
PIP_TIMEOUT="${PIP_TIMEOUT:-600}"
PIP_RETRIES="${PIP_RETRIES:-15}"
PIP_ATTEMPTS="${PIP_ATTEMPTS:-3}"

# 国内镜像（默认开启，海外服务器可 export USE_CN_MIRROR=0）
USE_CN_MIRROR="${USE_CN_MIRROR:-1}"
PYTORCH_INDEX_OFFICIAL="https://download.pytorch.org/whl/cu121"
PYTORCH_INDEX_CN="https://mirrors.tuna.tsinghua.edu.cn/pytorch-wheels/cu121"
PIP_INDEX_CN="https://pypi.tuna.tsinghua.edu.cn/simple"
HF_MIRROR="https://hf-mirror.com"

# ---------------------------------------------------------------------------
# 颜色输出
# ---------------------------------------------------------------------------
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
CYAN='\033[0;36m'
NC='\033[0m'

log_info()  { echo -e "${CYAN}[INFO]${NC}  $*"; }
log_ok()    { echo -e "${GREEN}[OK]${NC}    $*"; }
log_warn()  { echo -e "${YELLOW}[WARN]${NC}  $*"; }
log_error() { echo -e "${RED}[ERROR]${NC} $*" >&2; }

# ---------------------------------------------------------------------------
# 前置检查
# ---------------------------------------------------------------------------
require_root() {
    if [[ "${EUID:-$(id -u)}" -ne 0 ]]; then
        log_error "请使用 root 用户运行: sudo bash deploy.sh"
        exit 1
    fi
}

check_gpu() {
    if command -v nvidia-smi &>/dev/null; then
        log_ok "检测到 NVIDIA GPU:"
        nvidia-smi --query-gpu=name,driver_version,memory.total --format=csv,noheader
    else
        log_warn "未检测到 nvidia-smi，Whisper/ChatTTS CUDA 加速可能不可用。"
    fi
}

set_gpu_power_limit() {
    if command -v nvidia-smi &>/dev/null; then
        log_info "设置 GPU 功耗上限为 ${GPU_POWER_LIMIT}W ..."
        if nvidia-smi -pl "${GPU_POWER_LIMIT}" &>/dev/null; then
            log_ok "GPU 功耗墙已设为 ${GPU_POWER_LIMIT}W"
        else
            log_warn "无法设置功耗墙，请手动执行: nvidia-smi -pl ${GPU_POWER_LIMIT}"
        fi
    fi
}

# ---------------------------------------------------------------------------
# 系统依赖
# ---------------------------------------------------------------------------
install_system_deps() {
    log_info "安装系统依赖 ..."
    apt-get update -qq
    apt-get install -y \
        git curl wget build-essential \
        python3 python3-venv python3-dev python3-pip \
        ffmpeg libsndfile1 portaudio19-dev \
        ca-certificates gnupg

    log_ok "系统依赖安装完成"
}

install_node_pm2() {
    if command -v pm2 &>/dev/null; then
        log_ok "PM2 已安装: $(pm2 -v)"
        return
    fi

    log_info "安装 Node.js 20 LTS 与 PM2 ..."
    if ! command -v node &>/dev/null; then
        curl -fsSL https://deb.nodesource.com/setup_20.x | bash -
        apt-get install -y nodejs
    fi
    npm install -g pm2
    log_ok "PM2 安装完成: $(pm2 -v)"
}

# ---------------------------------------------------------------------------
# 代码部署
# ---------------------------------------------------------------------------
sync_git_repo() {
    local repo_dir="$1"
    cd "${repo_dir}"

    # 自动 stash 本地修改（如 sed 修 CRLF 等），避免 pull 被阻塞
    if [[ -n "$(git status --porcelain 2>/dev/null)" ]]; then
        log_warn "检测到本地未提交更改，自动 git stash ..."
        git stash push -u -m "deploy-auto-stash-$(date +%Y%m%d%H%M%S)" || true
    fi

    # 优先 ff-only pull
    if git pull --ff-only origin "${GIT_BRANCH}" 2>/dev/null; then
        log_ok "git pull 成功 (origin/${GIT_BRANCH})"
        return 0
    fi
    if git pull --ff-only 2>/dev/null; then
        log_ok "git pull 成功"
        return 0
    fi

    # pull 失败则强制与远端同步（生产服务器标准做法）
    log_warn "git pull 失败，执行 fetch + reset --hard 同步远端 ..."
    git fetch origin "${GIT_BRANCH}" 2>/dev/null || git fetch origin
    if git reset --hard "origin/${GIT_BRANCH}" 2>/dev/null; then
        log_ok "已强制同步到 origin/${GIT_BRANCH}"
        return 0
    fi
    if git reset --hard "origin/master" 2>/dev/null; then
        log_ok "已强制同步到 origin/master"
        return 0
    fi

    log_error "代码同步失败，请手动处理:"
    log_error "  cd ${repo_dir} && git status && git pull"
    return 1
}

deploy_code() {
    if [[ -d "${INSTALL_DIR}/.git" ]]; then
        log_info "更新已有代码: ${INSTALL_DIR}"
        sync_git_repo "${INSTALL_DIR}"
    elif [[ -d "${INSTALL_DIR}" ]]; then
        log_error "${INSTALL_DIR} 已存在但不是 git 仓库，请手动处理后重试。"
        exit 1
    else
        log_info "克隆仓库到 ${INSTALL_DIR} ..."
        git clone -b "${GIT_BRANCH}" "${GIT_REPO}" "${INSTALL_DIR}" || \
            git clone "${GIT_REPO}" "${INSTALL_DIR}"
    fi
    log_ok "代码就绪: ${INSTALL_DIR}"
}

# ---------------------------------------------------------------------------
# Python 环境
# ---------------------------------------------------------------------------
configure_pip() {
    local venv_pip="$1"

    log_info "配置 pip 网络参数 (timeout=${PIP_TIMEOUT}s, retries=${PIP_RETRIES}) ..."

    "${venv_pip}" config set global.timeout "${PIP_TIMEOUT}" 2>/dev/null || true
    "${venv_pip}" config set global.retries "${PIP_RETRIES}" 2>/dev/null || true

    if [[ "${USE_CN_MIRROR}" == "1" ]]; then
        log_info "启用国内镜像: 清华 PyPI + PyTorch cu121"
        "${venv_pip}" config set global.index-url "${PIP_INDEX_CN}" 2>/dev/null || true
        export PIP_INDEX_URL="${PIP_INDEX_CN}"
    fi

    # ChatTTS / HuggingFace 模型下载加速
    export HF_ENDPOINT="${HF_MIRROR}"
    export HF_HUB_ENABLE_HF_TRANSFER=0
}

pip_install_with_retry() {
    local attempt=1
    local pip_bin="$1"
    shift

    while [[ "${attempt}" -le "${PIP_ATTEMPTS}" ]]; do
        log_info "pip 安装中 (第 ${attempt}/${PIP_ATTEMPTS} 次)，大包下载较慢请耐心等待 ..."
        if "${pip_bin}" install \
            --timeout "${PIP_TIMEOUT}" \
            --retries "${PIP_RETRIES}" \
            --progress-bar on \
            "$@"; then
            return 0
        fi
        log_warn "pip 安装失败，60 秒后重试 ..."
        sleep 60
        attempt=$((attempt + 1))
    done

    log_error "pip 安装多次重试仍失败，请检查网络或手动安装后设置 SKIP_PYTORCH=1 重试"
    return 1
}

install_pytorch() {
    local pip_bin="$1"
    local pytorch_index="${PYTORCH_INDEX_OFFICIAL}"

    if [[ "${USE_CN_MIRROR}" == "1" ]]; then
        pytorch_index="${PYTORCH_INDEX_CN}"
    fi

    log_info "安装 PyTorch CUDA 12.1 (源: ${pytorch_index}) ..."
    log_warn "PyTorch + triton 体积约 2-3GB，国内网络可能需要 10-30 分钟，并非卡死。"

    # 分包安装，降低单次失败成本
    pip_install_with_retry "${pip_bin}" \
        torch \
        --index-url "${pytorch_index}"

    pip_install_with_retry "${pip_bin}" \
        torchvision torchaudio \
        --index-url "${pytorch_index}"
}

setup_python_venv() {
    local venv_path="${INSTALL_DIR}/venv"
    local pip_bin="${venv_path}/bin/pip"
    local python_bin="${venv_path}/bin/python"

    if [[ ! -d "${venv_path}" ]]; then
        log_info "创建 Python 虚拟环境 ..."
        python3 -m venv "${venv_path}"
    fi

    configure_pip "${pip_bin}"

    log_info "升级 pip ..."
    pip_install_with_retry "${pip_bin}" --upgrade pip setuptools wheel

    # 跳过已安装的 PyTorch
    if [[ "${SKIP_PYTORCH:-0}" == "1" ]]; then
        log_warn "SKIP_PYTORCH=1，跳过 PyTorch 安装"
    elif "${python_bin}" -c "import torch; assert torch.cuda.is_available()" 2>/dev/null; then
        log_ok "PyTorch 已安装且 CUDA 可用: $("${python_bin}" -c 'import torch; print(torch.__version__)')"
    else
        install_pytorch "${pip_bin}"
    fi

    log_info "安装项目依赖 (requirements.txt) ..."
    pip_install_with_retry "${pip_bin}" -r "${INSTALL_DIR}/requirements.txt"

    # 验证 CUDA
    if "${python_bin}" -c "import torch; assert torch.cuda.is_available()" 2>/dev/null; then
        log_ok "PyTorch CUDA 可用: $("${python_bin}" -c 'import torch; print(torch.cuda.get_device_name(0))')"
    else
        log_warn "PyTorch CUDA 不可用，请检查 NVIDIA 驱动与 CUDA 运行时。"
    fi

    log_ok "Python 虚拟环境配置完成"
}

create_runtime_dirs() {
    mkdir -p "${INSTALL_DIR}/logs"
    mkdir -p "${INSTALL_DIR}/uploads"
    mkdir -p "${INSTALL_DIR}/outputs"
    log_ok "运行时目录已创建 (logs, uploads, outputs)"
}

# ---------------------------------------------------------------------------
# 防火墙
# ---------------------------------------------------------------------------
configure_firewall() {
    if command -v ufw &>/dev/null && ufw status | grep -q "Status: active"; then
        log_info "放行 Gradio 端口 ${GRADIO_PORT} ..."
        ufw allow "${GRADIO_PORT}/tcp" || true
        log_ok "防火墙规则已更新"
    else
        log_info "ufw 未启用，跳过防火墙配置"
    fi
}

# ---------------------------------------------------------------------------
# PM2 管理
# ---------------------------------------------------------------------------
pm2_start() {
    log_info "通过 PM2 启动 Trading Studio ..."
    cd "${INSTALL_DIR}"

    if pm2 describe "${PM2_APP_NAME}" &>/dev/null; then
        pm2 delete "${PM2_APP_NAME}" || true
    fi

    pm2 start ecosystem.config.js
    pm2 save

    local startup_cmd
    startup_cmd=$(pm2 startup systemd -u root --hp /root 2>&1 | grep "sudo env" || true)
    if [[ -n "${startup_cmd}" ]]; then
        eval "${startup_cmd}" || log_warn "PM2 startup 可能已配置过"
    fi
    pm2 save

    log_ok "PM2 启动完成"
    pm2 status
}

pm2_restart() {
    cd "${INSTALL_DIR}"
    if pm2 describe "${PM2_APP_NAME}" &>/dev/null; then
        pm2 restart "${PM2_APP_NAME}"
        log_ok "PM2 已重启: ${PM2_APP_NAME}"
    else
        log_warn "进程不存在，执行完整启动 ..."
        pm2_start
    fi
    pm2 status
}

pm2_stop() {
    if pm2 describe "${PM2_APP_NAME}" &>/dev/null; then
        pm2 stop "${PM2_APP_NAME}"
        log_ok "PM2 已停止: ${PM2_APP_NAME}"
    else
        log_warn "PM2 进程 ${PM2_APP_NAME} 不存在"
    fi
    pm2 status
}

pm2_status() {
    echo ""
    log_info "=== PM2 状态 ==="
    pm2 status || true
    echo ""
    log_info "=== GPU 状态 ==="
    nvidia-smi 2>/dev/null || log_warn "nvidia-smi 不可用"
    echo ""
    log_info "=== 端口 ${GRADIO_PORT} 监听 ==="
    ss -tlnp | grep ":${GRADIO_PORT}" || log_warn "端口 ${GRADIO_PORT} 未监听，服务可能未启动"
    echo ""
    log_info "访问地址: http://$(hostname -I | awk '{print $1}'):${GRADIO_PORT}"
}

pm2_logs() {
    pm2 logs "${PM2_APP_NAME}" --lines 80 --nostream || true
}

# ---------------------------------------------------------------------------
# 主流程
# ---------------------------------------------------------------------------
cmd_install() {
    log_info "========== Trading Studio 一键部署开始 =========="
    log_info "安装目录: ${INSTALL_DIR}"
    log_info "运行用户: root"
    log_info "国内镜像: USE_CN_MIRROR=${USE_CN_MIRROR}  pip超时: ${PIP_TIMEOUT}s"

    install_system_deps
    install_node_pm2
    deploy_code
    setup_python_venv
    create_runtime_dirs
    set_gpu_power_limit
    configure_firewall
    pm2_start

    echo ""
    log_ok "========== 部署完成 =========="
    echo ""
    echo -e "  Web 中控: ${GREEN}http://$(hostname -I | awk '{print $1}'):${GRADIO_PORT}${NC}"
    echo -e "  项目目录: ${INSTALL_DIR}"
    echo -e "  查看日志: ${CYAN}pm2 logs ${PM2_APP_NAME}${NC}"
    echo -e "  重启服务: ${CYAN}bash ${INSTALL_DIR}/deploy.sh restart${NC}"
    echo ""
    log_warn "首次使用请打开 Web UI「音色锁定」上传参考人声，生成 speaker_emb.pt"
}

cmd_update() {
    log_info "========== 更新部署 =========="
    deploy_code
    setup_python_venv
    create_runtime_dirs
    pm2_restart
    log_ok "更新完成"
}

cmd_deps_only() {
    log_info "========== 仅安装/更新 Python 依赖 =========="
    setup_python_venv
    log_ok "依赖安装完成"
}

print_usage() {
    cat <<EOF
Trading Studio 一键部署脚本

用法:
  sudo bash deploy.sh [命令]

命令:
  (无参数)    首次完整部署到 /opt/Trading_Studio 并由 PM2 启动
  install     同上
  update      拉取最新代码、更新 Python 依赖、重启 PM2
  deps        仅安装/更新 Python 依赖（不拉代码、不启 PM2）
  restart     重启 PM2 进程
  stop        停止 PM2 进程
  status      查看 PM2 / GPU / 端口状态
  logs        查看 PM2 最近日志
  help        显示本帮助

环境变量:
  USE_CN_MIRROR=1   使用清华镜像（默认，国内推荐）
  USE_CN_MIRROR=0   使用 PyTorch 官方源
  PIP_TIMEOUT=600   pip 下载超时秒数
  SKIP_PYTORCH=1    跳过 PyTorch 安装

示例:
  cd /opt/Trading_Studio && bash deploy.sh update
  USE_CN_MIRROR=1 PIP_TIMEOUT=900 bash deploy.sh deps

EOF
}

main() {
    require_root

    local cmd="${1:-install}"

    case "${cmd}" in
        install|"")
            check_gpu
            cmd_install
            ;;
        update)
            check_gpu
            cmd_update
            ;;
        deps)
            check_gpu
            cmd_deps_only
            ;;
        restart)
            pm2_restart
            ;;
        stop)
            pm2_stop
            ;;
        status)
            pm2_status
            ;;
        logs)
            pm2_logs
            ;;
        help|-h|--help)
            print_usage
            ;;
        *)
            log_error "未知命令: ${cmd}"
            print_usage
            exit 1
            ;;
    esac
}

main "$@"