#!/bin/bash
# ========================================
# 智能矿山标签监测系统 - 一键升级脚本
# ========================================
# 使用方法:
#   bash upgrade.sh                          # 升级到最新master分支
#   bash upgrade.sh v1.2.0                   # 升级到指定标签版本
#   bash upgrade.sh abc123def                # 升级到指定commit
# ========================================

set -e

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 打印带颜色的消息
info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# ========================================
# 配置区域
# ========================================

# Git仓库地址（根据实际情况修改）
GIT_REPO="http://10.224.35.147:3000/redline/program.git"

# 服务名称（systemd）
SERVICE_NAME="mining-system"

# 需要保留的文件（不会被git覆盖）
PRESERVE_FILES=(
    "config.py"
    "mining_data.db"
    ".env"
)

# 需要保留的目录
PRESERVE_DIRS=(
    "outputs"
    "logs"
)

# ========================================
# 函数定义
# ========================================

# 获取项目根目录
get_project_dir() {
    echo "$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
}

# 检查是否安装了git
check_git() {
    if ! command -v git &> /dev/null; then
        error "Git未安装，请先安装: sudo apt-get install -y git"
        exit 1
    fi
}

# 检查当前是否在git仓库中
check_git_repo() {
    if ! git rev-parse --is-inside-work-tree &> /dev/null; then
        error "当前目录不是Git仓库"
        echo ""
        echo "如果需要首次使用Git升级方式，请先初始化："
        echo "  cd $PROJECT_DIR"
        echo "  git init"
        echo "  git remote add origin $GIT_REPO"
        echo "  git pull origin master"
        exit 1
    fi
}

# 获取当前版本信息
get_current_version() {
    if git rev-parse --is-inside-work-tree &> /dev/null; then
        local commit=$(git rev-parse --short HEAD 2>/dev/null || echo "unknown")
        local branch=$(git rev-parse --abbrev-ref HEAD 2>/dev/null || echo "unknown")
        local dirty=""
        if ! git diff --quiet &>/dev/null; then
            dirty=" (有未提交的修改)"
        fi
        echo "分支: $branch, 提交: $commit$dirty"
    else
        echo "未知（非Git仓库）"
    fi
}

# 备份当前系统
backup() {
    local DATE=$(date +%Y%m%d_%H%M%S)
    local BACKUP_DIR="${PROJECT_DIR}/backups"
    local BACKUP_FILE="${BACKUP_DIR}/backup-${DATE}.tar.gz"
    
    info "开始备份当前系统..."
    
    # 创建备份目录
    mkdir -p "$BACKUP_DIR"
    
    # 收集需要备份的文件
    local FILES_TO_BACKUP=()
    for file in "${PRESERVE_FILES[@]}"; do
        if [ -f "${PROJECT_DIR}/${file}" ]; then
            FILES_TO_BACKUP+=("$file")
        fi
    done
    
    # 备份配置文件和环境变量
    {
        echo "# 升级备份 - $DATE"
        echo "# 项目目录: $PROJECT_DIR"
        echo ""
        
        for file in "${FILES_TO_BACKUP[@]}"; do
            echo "===== $file ====="
            cat "${PROJECT_DIR}/${file}" 2>/dev/null || echo "(文件不存在)"
            echo ""
        done
        
        # 备份requirements.txt用于对比
        echo "===== requirements.txt ====="
        cat "${PROJECT_DIR}/requirements.txt" 2>/dev/null || echo "(文件不存在)"
        echo ""
        
        # 备份当前git信息
        echo "===== git info ====="
        echo "branch: $(git rev-parse --abbrev-ref HEAD 2>/dev/null || echo 'N/A')"
        echo "commit: $(git rev-parse HEAD 2>/dev/null || echo 'N/A')"
        echo ""
    } > "${BACKUP_DIR}/config-backup-${DATE}.txt"
    
    # 复制需要保留的文件到备份目录
    for file in "${PRESERVE_FILES[@]}"; do
        if [ -f "${PROJECT_DIR}/${file}" ]; then
            cp "${PROJECT_DIR}/${file}" "${BACKUP_DIR}/${file}.${DATE}.bak"
        fi
    done
    
    # 备份数据库（单独备份，确保完整）
    if [ -f "${PROJECT_DIR}/mining_data.db" ]; then
        cp "${PROJECT_DIR}/mining_data.db" "${BACKUP_DIR}/mining_data.db.${DATE}.bak"
        info "数据库已备份: mining_data.db.${DATE}.bak"
    fi
    
    # 保留当前outputs目录
    if [ -d "${PROJECT_DIR}/outputs" ]; then
        tar -czf "${BACKUP_DIR}/outputs-${DATE}.tar.gz" -C "$PROJECT_DIR" outputs 2>/dev/null || true
    fi
    
    # 保留最近的10个备份，清理旧备份
    cd "$BACKUP_DIR"
    ls -t *.tar.gz 2>/dev/null | tail -n +11 | xargs -r rm -f
    ls -t *.txt 2>/dev/null | tail -n +11 | xargs -r rm -f
    ls -t *.bak 2>/dev/null | tail -n +20 | xargs -r rm -f
    cd "$PROJECT_DIR"
    
    success "备份完成: ${BACKUP_DIR}/"
    echo ""
}

# 停止服务
stop_service() {
    info "停止服务..."
    
    # 尝试systemd方式
    if systemctl list-unit-files | grep -q "${SERVICE_NAME}.service" 2>/dev/null; then
        sudo systemctl stop "$SERVICE_NAME" 2>/dev/null && \
            success "已通过systemd停止服务" && return
    fi
    
    # 尝试kill端口进程
    if command -v lsof &> /dev/null; then
        local PIDS=$(lsof -ti:5000 2>/dev/null || true)
        if [ -n "$PIDS" ]; then
            echo "$PIDS" | xargs kill -TERM 2>/dev/null || true
            sleep 2
            echo "$PIDS" | xargs kill -KILL 2>/dev/null || true
            success "已停止端口5000上的进程"
            return
        fi
    fi
    
    # 尝试ps方式
    local PIDS=$(ps aux | grep '[p]ython.*app.py' | awk '{print $2}' || true)
    if [ -n "$PIDS" ]; then
        echo "$PIDS" | xargs kill -TERM 2>/dev/null || true
        sleep 2
        echo "$PIDS" | xargs kill -KILL 2>/dev/null || true
        success "已停止Python应用进程"
        return
    fi
    
    warn "未发现运行中的服务进程"
    echo ""
}

# 保存需要保留的文件
save_preserved_files() {
    info "保存需要保留的文件..."
    
    local STAGING_DIR="${PROJECT_DIR}/.upgrade_staging"
    mkdir -p "$STAGING_DIR"
    
    # 保存文件
    for file in "${PRESERVE_FILES[@]}"; do
        if [ -f "${PROJECT_DIR}/${file}" ]; then
            cp -f "${PROJECT_DIR}/${file}" "${STAGING_DIR}/${file}"
            info "  已保存: $file"
        fi
    done
    
    # 保存目录内容
    for dir in "${PRESERVE_DIRS[@]}"; do
        if [ -d "${PROJECT_DIR}/${dir}" ]; then
            # 先删除staging中的旧目录（如果有），再复制
            rm -rf "${STAGING_DIR}/${dir}" 2>/dev/null || true
            cp -r "${PROJECT_DIR}/${dir}" "${STAGING_DIR}/${dir}"
            info "  已保存目录: $dir/"
        fi
    done
    
    echo ""
}

# 从Git拉取最新代码
pull_latest_code() {
    local TARGET="${1:-master}"
    
    info "从Git拉取最新代码 (目标: $TARGET)..."
    
    # 检查远程仓库
    if ! git remote -v | grep -q "origin" 2>/dev/null; then
        error "未找到Git远程仓库 'origin'"
        echo "添加远程仓库: git remote add origin $GIT_REPO"
        git remote add origin "$GIT_REPO" 2>/dev/null || true
    fi
    
    # 如果远程地址不正确，更新它
    local CURRENT_REMOTE=$(git remote get-url origin 2>/dev/null || echo "")
    if [ "$CURRENT_REMOTE" != "$GIT_REPO" ] && [ -n "$CURRENT_REMOTE" ]; then
        warn "远程地址已更新: $CURRENT_REMOTE -> $GIT_REPO"
        git remote set-url origin "$GIT_REPO"
    fi
    
    # 拉取最新代码
    # 先stash本地修改（已保存到staging，所以可以强制）
    git stash --include-untracked 2>/dev/null || true
    
    # 根据目标拉取
    if [[ "$TARGET" == *"."* ]] || [[ "$TARGET" == v* ]]; then
        # 可能是tag
        git fetch origin --tags 2>&1 | grep -v "already up" || true
        if git tag -l | grep -q "^${TARGET}\$"; then
            info "切换到标签: $TARGET"
            git checkout tags/"$TARGET" -- . 2>&1 || git pull origin "$TARGET" 2>&1 || true
        else
            info "拉取master分支最新代码"
            git checkout master 2>&1 || true
            git pull origin master 2>&1 || true
        fi
    elif [[ "$TARGET" =~ ^[0-9a-f]{7,40}$ ]]; then
        # 可能是commit hash
        info "切换到commit: $TARGET"
        git fetch origin 2>&1 | grep -v "already up" || true
        git checkout "$TARGET" -- . 2>&1 || git pull origin master 2>&1 || true
    else
        # 分支名
        info "拉取分支: $TARGET"
        git checkout "$TARGET" 2>&1 || git checkout master 2>&1 || true
        git pull origin "$TARGET" 2>&1 || git pull origin master 2>&1 || true
    fi
    
    success "代码拉取完成"
    echo ""
}

# 恢复保留的文件
restore_preserved_files() {
    info "恢复保留的文件..."
    
    local STAGING_DIR="${PROJECT_DIR}/.upgrade_staging"
    
    if [ ! -d "$STAGING_DIR" ]; then
        warn "暂存目录不存在，跳过恢复"
        return
    fi
    
    # 恢复文件
    for file in "${PRESERVE_FILES[@]}"; do
        if [ -f "${STAGING_DIR}/${file}" ]; then
            cp -f "${STAGING_DIR}/${file}" "${PROJECT_DIR}/${file}"
            info "  已恢复: $file"
        fi
    done
    
    # 恢复目录内容（合并方式）
    for dir in "${PRESERVE_DIRS[@]}"; do
        if [ -d "${STAGING_DIR}/${dir}" ]; then
            mkdir -p "${PROJECT_DIR}/${dir}"
            rsync -a "${STAGING_DIR}/${dir}/" "${PROJECT_DIR}/${dir}/" 2>/dev/null || \
                cp -r "${STAGING_DIR}/${dir}"/* "${PROJECT_DIR}/${dir}/" 2>/dev/null || true
            info "  已恢复目录: $dir/"
        fi
    done
    
    # 清理暂存目录
    rm -rf "$STAGING_DIR"
    
    echo ""
}

# 更新Python依赖
update_dependencies() {
    info "检查并更新Python依赖..."
    
    local VENV_DIR="${PROJECT_DIR}/venv"
    
    # 检查虚拟环境是否存在
    if [ ! -d "$VENV_DIR" ]; then
        info "创建Python虚拟环境..."
        python3 -m venv "$VENV_DIR"
    fi
    
    # 激活虚拟环境
    source "${VENV_DIR}/bin/activate"
    
    # 升级pip
    pip install --upgrade pip -q 2>&1 | grep -v "already satisfied" || true
    
    # 检查requirements.txt是否有变化
    if [ -f "${PROJECT_DIR}/requirements.txt" ]; then
        pip install -r "${PROJECT_DIR}/requirements.txt" 2>&1 | grep -v "already satisfied" || true
        success "依赖更新完成"
    else
        error "requirements.txt 不存在！"
        exit 1
    fi
    
    # 退出虚拟环境（不影响后续命令）
    deactivate 2>/dev/null || true
    
    echo ""
}

# 运行数据库迁移（如有需要）
run_migrations() {
    info "检查数据库迁移..."
    
    # 如果有migration目录，执行迁移
    if [ -d "${PROJECT_DIR}/migrations" ]; then
        info "发现迁移目录，执行迁移..."
        source "${PROJECT_DIR}/venv/bin/activate"
        cd "${PROJECT_DIR}/migrations"
        python3 migrate.py 2>/dev/null || warn "迁移脚本执行失败或无需迁移"
        cd "$PROJECT_DIR"
        deactivate 2>/dev/null || true
    else
        success "无需数据库迁移"
    fi
    
    echo ""
}

# 重启服务
restart_service() {
    info "重启服务..."
    
    # 尝试systemd方式
    if systemctl list-unit-files | grep -q "${SERVICE_NAME}.service" 2>/dev/null; then
        sudo systemctl daemon-reload
        sudo systemctl restart "$SERVICE_NAME"
        sleep 2
        if sudo systemctl is-active --quiet "$SERVICE_NAME"; then
            success "已通过systemd重启服务"
        else
            error "systemd服务启动失败，请查看日志: sudo journalctl -u $SERVICE_NAME -f"
        fi
        return
    fi
    
    # 如果没有systemd服务，提示手动启动
    warn "未检测到systemd服务，请手动启动："
    echo ""
    echo "  cd $PROJECT_DIR"
    echo "  source venv/bin/activate"
    echo "  python3 app.py"
    echo ""
    echo "或使用启动脚本:"
    echo "  bash $PROJECT_DIR/start.sh"
    echo ""
}

# 验证升级结果
verify_upgrade() {
    info "验证升级结果..."
    
    sleep 3
    
    # 检查API是否可访问
    if command -v curl &> /dev/null; then
        local RESPONSE=$(curl -s -o /dev/null -w "%{http_code}" http://localhost:5000/ 2>/dev/null || echo "000")
        if [ "$RESPONSE" = "200" ] || [ "$RESPONSE" = "302" ]; then
            success "API响应正常 (HTTP $RESPONSE)"
        else
            warn "API响应异常 (HTTP $RESPONSE)，请检查服务状态"
        fi
        
        # 检查stats API
        local STATS_RESPONSE=$(curl -s http://localhost:5000/api/stats/all 2>/dev/null || echo "")
        if echo "$STATS_RESPONSE" | python3 -m json.tool &>/dev/null; then
            success "统计API正常"
        fi
    fi
    
    # 显示新版本信息
    echo ""
    info "新版本信息:"
    if git rev-parse --is-inside-work-tree &> /dev/null; then
        echo "  分支: $(git rev-parse --abbrev-ref HEAD 2>/dev/null || echo 'unknown')"
        echo "  提交: $(git rev-parse --short HEAD 2>/dev/null || echo 'unknown')"
        echo "  时间: $(git log -1 --format='%ci' 2>/dev/null || echo 'unknown')"
        echo "  信息: $(git log -1 --format='%s' 2>/dev/null || echo 'unknown')"
    fi
    
    echo ""
}

# 显示回滚说明
show_rollback_info() {
    local BACKUP_DIR="${PROJECT_DIR}/backups"
    
    echo "======================================"
    echo "  回滚方案（如需）"
    echo "======================================"
    echo ""
    echo "如果升级后发现问题，可以执行回滚："
    echo ""
    echo "  # 方式1: 使用升级脚本回滚到上一个备份"
    echo "  bash $PROJECT_DIR/upgrade.sh rollback"
    echo ""
    echo "  # 方式2: 手动回滚"
    echo "  cd $PROJECT_DIR/backups"
    echo "  ls -lt  # 找到最新的备份"
    echo "  # 恢复配置文件"
    echo "  cp config.py.YYYYMMDD_HHMMSS.bak $PROJECT_DIR/"
    echo "  # 恢复数据库"
    echo "  cp mining_data.db.YYYYMMDD_HHMMSS.bak $PROJECT_DIR/"
    echo "  # 回滚git代码"
    echo "  git checkout <旧版本commit>"
    echo "  # 重启服务"
    echo "  sudo systemctl restart $SERVICE_NAME"
    echo ""
}

# 回滚到上一个备份
rollback() {
    local BACKUP_DIR="${PROJECT_DIR}/backups"
    
    if [ ! -d "$BACKUP_DIR" ]; then
        error "备份目录不存在，无法回滚"
        exit 1
    fi
    
    # 找到最新的配置备份
    local LATEST_CONFIG=$(ls -t "${BACKUP_DIR}"/config-backup-*.txt 2>/dev/null | head -1)
    if [ -z "$LATEST_CONFIG" ]; then
        error "未找到配置备份文件"
        exit 1
    fi
    
    info "找到备份: $LATEST_CONFIG"
    
    # 询问确认
    echo ""
    echo "备份内容预览:"
    head -5 "$LATEST_CONFIG"
    echo "..."
    echo ""
    read -p "确认回滚到此备份? (y/N) " -n 1 -r
    echo
    if [[ ! $REPLY =~ ^[Yy]$ ]]; then
        info "回滚已取消"
        exit 0
    fi
    
    # 停止服务
    stop_service
    
    # 提取备份中的日期
    local BACKUP_DATE=$(basename "$LATEST_CONFIG" | sed 's/config-backup-//;s/.txt//')
    
    # 恢复配置文件
    for file in "${PRESERVE_FILES[@]}"; do
        local BAK_FILE="${BACKUP_DIR}/${file}.${BACKUP_DATE}.bak"
        if [ -f "$BAK_FILE" ]; then
            cp -f "$BAK_FILE" "${PROJECT_DIR}/${file}"
            info "已恢复: $file"
        fi
    done
    
    # 恢复outputs目录
    local OUTPUTS_BACKUP=$(ls -t "${BACKUP_DIR}"/outputs-*.tar.gz 2>/dev/null | head -1)
    if [ -n "$OUTPUTS_BACKUP" ]; then
        tar -xzf "$OUTPUTS_BACKUP" -C "$PROJECT_DIR" 2>/dev/null && \
            info "已恢复: outputs/"
    fi
    
    # 回滚git（尝试回到上一个commit）
    if git rev-parse --is-inside-work-tree &> /dev/null; then
        # 从备份文件中获取之前的commit
        local OLD_COMMIT=$(grep "^commit:" "$LATEST_CONFIG" | awk '{print $2}' | head -1)
        if [ -n "$OLD_COMMIT" ] && [ "$OLD_COMMIT" != "N/A" ]; then
            info "回滚Git到commit: $OLD_COMMIT"
            git checkout "$OLD_COMMIT" 2>/dev/null || \
                git reset --hard "$OLD_COMMIT" 2>/dev/null || \
                warn "Git回滚失败，请手动操作"
        fi
    fi
    
    # 更新依赖
    update_dependencies
    
    # 重启服务
    restart_service
    
    success "回滚完成"
}

# 显示帮助
show_help() {
    echo "======================================"
    echo "  智能矿山标签监测系统 - 升级脚本"
    echo "======================================"
    echo ""
    echo "使用方法:"
    echo "  bash upgrade.sh [选项]"
    echo ""
    echo "选项:"
    echo "  (无参数)          升级到master分支最新版本"
    echo "  <version>         升级到指定版本(tag/commit/分支)"
    echo "  rollback          回滚到上一个备份"
    echo "  status            显示当前版本状态"
    echo "  backup            仅执行备份"
    echo "  help              显示此帮助信息"
    echo ""
    echo "示例:"
    echo "  bash upgrade.sh               # 升级到最新"
    echo "  bash upgrade.sh v1.2.0        # 升级到v1.2.0标签"
    echo "  bash upgrade.sh abc123def     # 升级到指定commit"
    echo "  bash upgrade.sh develop       # 切换到develop分支"
    echo "  bash upgrade.sh rollback      # 回滚"
    echo "  bash upgrade.sh status        # 查看版本状态"
    echo ""
    echo "配置:"
    echo "  Git仓库: $GIT_REPO"
    echo "  服务名:  $SERVICE_NAME"
    echo ""
    echo "保留文件（升级时不覆盖）:"
    for file in "${PRESERVE_FILES[@]}"; do
        echo "  - $file"
    done
    echo ""
    echo "保留目录:"
    for dir in "${PRESERVE_DIRS[@]}"; do
        echo "  - $dir/"
    done
    echo ""
}

# 显示当前状态
show_status() {
    echo "======================================"
    echo "  当前版本状态"
    echo "======================================"
    echo ""
    echo "项目目录: $PROJECT_DIR"
    echo ""
    
    if git rev-parse --is-inside-work-tree &> /dev/null; then
        echo "Git信息:"
        echo "  分支:   $(git rev-parse --abbrev-ref HEAD 2>/dev/null || echo 'unknown')"
        echo "  提交:   $(git rev-parse --short HEAD 2>/dev/null || echo 'unknown')"
        echo "  完整哈希: $(git rev-parse HEAD 2>/dev/null || echo 'unknown')"
        echo "  时间:   $(git log -1 --format='%ci' 2>/dev/null || echo 'unknown')"
        echo "  信息:   $(git log -1 --format='%s' 2>/dev/null || echo 'unknown')"
        echo "  远程:   $(git remote get-url origin 2>/dev/null || echo 'unknown')"
        
        # 检查是否有更新
        echo ""
        info "检查远程仓库更新..."
        git fetch origin 2>&1 | grep -v "already up" || true
        local LOCAL=$(git rev-parse HEAD 2>/dev/null)
        local REMOTE=$(git rev-parse @{upstream} 2>/dev/null || echo "")
        if [ "$LOCAL" = "$REMOTE" ]; then
            success "已是最新版本"
        elif [ -n "$REMOTE" ]; then
            local AHEAD=$(git rev-list_count ..@{upstream} 2>/dev/null || echo "?")
            warn "有可用更新"
        fi
        
        # 检查工作区状态
        echo ""
        if git diff --quiet &>/dev/null; then
            echo "工作区: 干净"
        else
            warn "工作区有未提交的修改"
            git status --short 2>/dev/null | head -10
        fi
    else
        echo "非Git仓库"
    fi
    
    echo ""
    echo "服务状态:"
    if systemctl is-active --quiet "$SERVICE_NAME" 2>/dev/null; then
        echo "  systemd服务: 运行中"
    else
        echo "  systemd服务: 未运行或未配置"
    fi
    
    # 检查进程
    local PIDS=$(ps aux | grep '[p]ython.*app.py' | wc -l || echo "0")
    echo "  Python进程: $PIDS"
    
    echo ""
    echo "备份信息:"
    local BACKUP_DIR="${PROJECT_DIR}/backups"
    if [ -d "$BACKUP_DIR" ]; then
        local BACKUP_COUNT=$(ls -1 "$BACKUP_DIR" 2>/dev/null | wc -l)
        echo "  备份数量: $BACKUP_COUNT"
        echo "  最新备份: $(ls -t "$BACKUP_DIR" 2>/dev/null | head -1)"
    else
        echo "  暂无备份"
    fi
    
    echo ""
}

# ========================================
# 主流程
# ========================================

# 获取项目目录
PROJECT_DIR=$(get_project_dir)

echo "======================================"
echo "  智能矿山标签监测系统"
echo "  一键升级脚本"
echo "======================================"
echo ""
echo "项目目录: $PROJECT_DIR"
echo ""

# 处理特殊命令
case "${1:-}" in
    rollback)
        check_git
        check_git_repo
        rollback
        exit 0
        ;;
    status)
        check_git
        check_git_repo
        show_status
        exit 0
        ;;
    backup)
        check_git
        check_git_repo
        backup
        exit 0
        ;;
    help|--help|-h)
        show_help
        exit 0
        ;;
esac

# 前置检查
check_git
check_git_repo

# 显示当前版本
echo "当前版本: $(get_current_version)"
TARGET="${1:-master}"
echo "目标版本: $TARGET"
echo ""

# 步骤1: 备份
backup

# 步骤2: 停止服务
stop_service

# 步骤3: 保存需要保留的文件
save_preserved_files

# 步骤4: 拉取最新代码
pull_latest_code "$TARGET"

# 步骤5: 恢复保留的文件
restore_preserved_files

# 步骤6: 更新依赖
update_dependencies

# 步骤7: 数据库迁移
run_migrations

# 步骤8: 重启服务
restart_service

# 步骤9: 验证
verify_upgrade

# 完成
echo "======================================"
echo -e "${GREEN}  升级完成!${NC}"
echo "======================================"
echo ""

# 显示回滚说明
show_rollback_info