#!/bin/bash
# TSP智能助手升级脚本

set -e

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

# 日志函数
log_info() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

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

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

log_step() {
    echo -e "${BLUE}[STEP]${NC} $1"
}

# 配置变量
APP_NAME="tsp_assistant"
BACKUP_DIR="./backups"
DEPLOY_PATH="/opt/tsp_assistant"
SERVICE_NAME="tsp_assistant"
HEALTH_URL="http://localhost:5000/api/health"

# 检查参数
if [ $# -lt 1 ]; then
    echo "用法: $0 <新版本路径> [选项]"
    echo "选项:"
    echo "  --force    强制升级，跳过确认"
    echo "  --no-backup 跳过备份"
    echo "  --rollback  回滚到指定备份"
    exit 1
fi

NEW_VERSION_PATH=$1
FORCE_UPGRADE=false
SKIP_BACKUP=false
ROLLBACK_MODE=false

# 解析参数
while [[ $# -gt 1 ]]; do
    case $2 in
        --force)
            FORCE_UPGRADE=true
            ;;
        --no-backup)
            SKIP_BACKUP=true
            ;;
        --rollback)
            ROLLBACK_MODE=true
            ;;
        *)
            log_error "未知选项: $2"
            exit 1
            ;;
    esac
    shift
done

# 回滚功能
rollback() {
    local backup_name=$1
    
    if [ -z "$backup_name" ]; then
        log_error "请指定备份名称"
        exit 1
    fi
    
    log_step "开始回滚到备份: $backup_name"
    
    # 检查备份是否存在
    if [ ! -d "$BACKUP_DIR/$backup_name" ]; then
        log_error "备份不存在: $backup_name"
        log_info "可用备份列表:"
        ls -la "$BACKUP_DIR" | grep backup
        exit 1
    fi
    
    # 停止服务
    log_info "停止服务..."
    sudo systemctl stop "$SERVICE_NAME" || true
    
    # 恢复文件
    log_info "恢复文件..."
    sudo rm -rf "$DEPLOY_PATH"
    sudo cp -r "$BACKUP_DIR/$backup_name" "$DEPLOY_PATH"
    sudo chown -R www-data:www-data "$DEPLOY_PATH"
    
    # 恢复数据库
    if [ -f "$BACKUP_DIR/$backup_name/database/tsp_assistant.db" ]; then
        log_info "恢复数据库..."
        sudo cp "$BACKUP_DIR/$backup_name/database/tsp_assistant.db" "$DEPLOY_PATH/"
    fi
    
    # 启动服务
    log_info "启动服务..."
    sudo systemctl start "$SERVICE_NAME"
    
    # 等待服务启动
    sleep 10
    
    # 健康检查
    if curl -f "$HEALTH_URL" > /dev/null 2>&1; then
        log_info "回滚成功！"
    else
        log_error "回滚后健康检查失败"
        exit 1
    fi
}

# 创建备份
create_backup() {
    local timestamp=$(date +"%Y%m%d_%H%M%S")
    local backup_name="${APP_NAME}_backup_${timestamp}"
    local backup_path="$BACKUP_DIR/$backup_name"
    
    log_step "创建备份: $backup_name"
    
    # 创建备份目录
    mkdir -p "$backup_path"
    
    # 备份应用文件
    if [ -d "$DEPLOY_PATH" ]; then
        log_info "备份应用文件..."
        cp -r "$DEPLOY_PATH"/* "$backup_path/"
    fi
    
    # 备份数据库
    if [ -f "$DEPLOY_PATH/tsp_assistant.db" ]; then
        log_info "备份数据库..."
        mkdir -p "$backup_path/database"
        cp "$DEPLOY_PATH/tsp_assistant.db" "$backup_path/database/"
    fi
    
    # 保存备份信息
    cat > "$backup_path/backup_info.json" << EOF
{
  "backup_name": "$backup_name",
  "backup_path": "$backup_path",
  "timestamp": "$timestamp",
  "version": "$(cd "$DEPLOY_PATH" && python version.py version 2>/dev/null || echo "unknown")",
  "git_commit": "$(cd "$DEPLOY_PATH" && git rev-parse HEAD 2>/dev/null | cut -c1-8 || echo "unknown")"
}
EOF
    
    log_info "备份完成: $backup_name"
    echo "$backup_name"
}

# 升级功能
upgrade() {
    local new_version_path=$1
    
    log_step "开始升级TSP智能助手"
    
    # 检查新版本路径
    if [ ! -d "$new_version_path" ]; then
        log_error "新版本路径不存在: $new_version_path"
        exit 1
    fi
    
    # 检查当前版本
    if [ -d "$DEPLOY_PATH" ]; then
        local current_version=$(cd "$DEPLOY_PATH" && python version.py version 2>/dev/null || echo "unknown")
        log_info "当前版本: $current_version"
    else
        log_warn "当前部署路径不存在: $DEPLOY_PATH"
    fi
    
    # 检查新版本
    local new_version=$(cd "$new_version_path" && python version.py version 2>/dev/null || echo "unknown")
    log_info "新版本: $new_version"
    
    # 确认升级
    if [ "$FORCE_UPGRADE" = false ]; then
        echo -n "确认升级到版本 $new_version? (y/N): "
        read -r response
        if [[ ! "$response" =~ ^[Yy]$ ]]; then
            log_info "升级取消"
            exit 0
        fi
    fi
    
    # 创建备份
    local backup_name=""
    if [ "$SKIP_BACKUP" = false ]; then
        backup_name=$(create_backup)
    fi
    
    # 停止服务
    log_step "停止服务..."
    sudo systemctl stop "$SERVICE_NAME" || true
    
    # 升级文件
    log_step "升级应用文件..."
    sudo rm -rf "$DEPLOY_PATH"
    sudo mkdir -p "$DEPLOY_PATH"
    sudo cp -r "$new_version_path"/* "$DEPLOY_PATH/"
    sudo chown -R www-data:www-data "$DEPLOY_PATH"
    
    # 安装依赖
    log_step "安装依赖..."
    cd "$DEPLOY_PATH"
    sudo -u www-data python -m pip install -r requirements.txt
    
    # 运行数据库迁移
    log_step "运行数据库迁移..."
    sudo -u www-data python init_database.py || true
    
    # 启动服务
    log_step "启动服务..."
    sudo systemctl start "$SERVICE_NAME"
    
    # 等待服务启动
    log_info "等待服务启动..."
    sleep 15
    
    # 健康检查
    log_step "执行健康检查..."
    local retry_count=0
    local max_retries=10
    
    while [ $retry_count -lt $max_retries ]; do
        if curl -f "$HEALTH_URL" > /dev/null 2>&1; then
            log_info "健康检查通过！"
            break
        else
            log_warn "健康检查失败，重试中... ($((retry_count + 1))/$max_retries)"
            retry_count=$((retry_count + 1))
            sleep 5
        fi
    done
    
    if [ $retry_count -eq $max_retries ]; then
        log_error "健康检查失败，开始回滚..."
        if [ -n "$backup_name" ]; then
            rollback "$backup_name"
        else
            log_error "没有备份可回滚"
            exit 1
        fi
    else
        log_info "升级成功！"
        log_info "新版本: $new_version"
        if [ -n "$backup_name" ]; then
            log_info "备份名称: $backup_name"
        fi
    fi
}

# 主函数
main() {
    if [ "$ROLLBACK_MODE" = true ]; then
        rollback "$NEW_VERSION_PATH"
    else
        upgrade "$NEW_VERSION_PATH"
    fi
}

# 执行主函数
main