#!/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 最近日志
# =============================================================================

set -euo pipefail

# ---------------------------------------------------------------------------
# 可配置常量
# ---------------------------------------------------------------------------
INSTALL_DIR="/opt/Trading_Studio"
GIT_REPO="https://git.bz121.com/dekun/Trading_Studio.git"
PM2_APP_NAME="trading_studio"
GRADIO_PORT=5683
GPU_POWER_LIMIT=120
PYTORCH_INDEX="https://download.pytorch.org/whl/cu121"

# ---------------------------------------------------------------------------
# 颜色输出
# ---------------------------------------------------------------------------
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)"
}

# ---------------------------------------------------------------------------
# 代码部署
# ---------------------------------------------------------------------------
deploy_code() {
    if [[ -d "${INSTALL_DIR}/.git" ]]; then
        log_info "更新已有代码: ${INSTALL_DIR}"
        git -C "${INSTALL_DIR}" pull --ff-only || {
            log_warn "git pull 失败，尝试保留本地更改继续部署 ..."
        }
    elif [[ -d "${INSTALL_DIR}" ]]; then
        log_error "${INSTALL_DIR} 已存在但不是 git 仓库，请手动处理后重试。"
        exit 1
    else
        log_info "克隆仓库到 ${INSTALL_DIR} ..."
        git clone "${GIT_REPO}" "${INSTALL_DIR}"
    fi
    log_ok "代码就绪: ${INSTALL_DIR}"
}

# ---------------------------------------------------------------------------
# Python 环境
# -------------------------------------------------------------------
setup_python_venv() {
    local venv_path="${INSTALL_DIR}/venv"

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

    # shellcheck disable=SC1091
    source "${venv_path}/bin/activate"

    log_info "升级 pip ..."
    pip install --upgrade pip setuptools wheel -q

    log_info "安装 PyTorch (CUDA 12.1) ..."
    pip install torch torchvision torchaudio --index-url "${PYTORCH_INDEX}" -q

    log_info "安装项目依赖 ..."
    pip install -r "${INSTALL_DIR}/requirements.txt" -q

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

    deactivate
    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

    # 配置 root 用户开机自启
    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"

    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 "更新完成"
}

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

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

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

示例:
  cd /opt/Trading_Studio && sudo bash deploy.sh update

EOF
}

main() {
    require_root

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

    case "${cmd}" in
        install|"")
            check_gpu
            cmd_install
            ;;
        update)
            check_gpu
            cmd_update
            ;;
        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 "$@"