#!/bin/bash

# ========================================
# 漏洞管理后端服务管理脚本 (macOS)
# ========================================

set -e  # 遇到错误立即退出

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

# 项目根目录
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)"

# 配置
MYSQL_CONTAINER_NAME="vul-backend-mysql"
MYSQL_ROOT_PASSWORD="securitycenter"
MYSQL_DATABASE="vuln_management"
MYSQL_PORT="3306"
API_PORT="8080"

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

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

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

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

# 打印使用说明
print_usage() {
    echo ""
    echo -e "${CYAN}========================================${NC}"
    echo -e "${CYAN}   漏洞管理后端服务管理脚本${NC}"
    echo -e "${CYAN}========================================${NC}"
    echo ""
    echo -e "${GREEN}使用方法:${NC}"
    echo -e "  $0 ${YELLOW}start${NC}    - 启动所有服务（MySQL + 后端）"
    echo -e "  $0 ${YELLOW}stop${NC}     - 停止所有服务"
    echo -e "  $0 ${YELLOW}restart${NC}  - 重启所有服务"
    echo -e "  $0 ${YELLOW}status${NC}   - 查看服务状态"
    echo -e "  $0 ${YELLOW}logs${NC}     - 查看服务日志"
    echo -e "  $0 ${YELLOW}clean${NC}    - 停止并清理所有数据（危险操作）"
    echo ""
    echo -e "${GREEN}示例:${NC}"
    echo -e "  $0 start     # 一键启动所有服务"
    echo -e "  $0 stop      # 停止所有服务"
    echo -e "  $0 restart   # 重启服务"
    echo ""
}

# 打印标题
print_header() {
    echo ""
    echo -e "${GREEN}========================================${NC}"
    echo -e "${GREEN}   $1${NC}"
    echo -e "${GREEN}========================================${NC}"
    echo ""
}

# 检查命令是否存在
check_command() {
    if ! command -v "$1" &> /dev/null; then
        log_error "$1 未安装，请先安装 $1"
        exit 1
    fi
}

# 检查必要的依赖
check_dependencies() {
    log_info "检查系统依赖..."
    
    check_command "docker"
    check_command "docker-compose"
    check_command "go"
    
    # 检查 Docker 是否运行
    if ! docker info &> /dev/null; then
        log_error "Docker 未运行，请先启动 Docker Desktop"
        exit 1
    fi
    
    log_success "系统依赖检查通过"
}

# 检查端口是否被占用
check_port() {
    local port=$1
    if lsof -Pi :$port -sTCP:LISTEN -t &> /dev/null; then
        log_warning "端口 $port 已被占用"
        local pid=$(lsof -ti:$port)
        log_warning "占用进程 PID: $pid"
        
        read -p "是否终止占用端口的进程? (y/n): " -n 1 -r
        echo
        if [[ $REPLY =~ ^[Yy]$ ]]; then
            kill -9 $pid
            log_success "已终止进程 $pid"
            sleep 1
        else
            log_error "端口 $port 被占用，无法继续启动"
            exit 1
        fi
    fi
}

# 停止并清理旧容器
cleanup_old_containers() {
    log_info "检查并清理旧容器..."
    
    if docker ps -a | grep -q "$MYSQL_CONTAINER_NAME"; then
        log_info "发现旧的 MySQL 容器，正在清理..."
        docker stop "$MYSQL_CONTAINER_NAME" 2>/dev/null || true
        docker rm "$MYSQL_CONTAINER_NAME" 2>/dev/null || true
        log_success "旧容器已清理"
    else
        log_info "没有发现旧容器"
    fi
}

# 启动 MySQL 容器
start_mysql() {
    log_info "启动 MySQL 容器..."
    
    cd "$SCRIPT_DIR"
    
    # 检查 MySQL 端口
    check_port "$MYSQL_PORT"
    
    # 使用 docker-compose 启动
    docker-compose up -d
    
    log_info "等待 MySQL 启动..."
    
    # 等待 MySQL 就绪（最多等待 60 秒）
    local max_attempts=60
    local attempt=0
    
    while [ $attempt -lt $max_attempts ]; do
        if docker exec "$MYSQL_CONTAINER_NAME" mysqladmin ping -h localhost -uroot -p"$MYSQL_ROOT_PASSWORD" &> /dev/null; then
            log_success "MySQL 容器已启动并就绪"
            return 0
        fi
        
        echo -n "."
        sleep 1
        ((attempt++))
    done
    
    log_error "MySQL 启动超时"
    exit 1
}

# 初始化数据库
init_database() {
    log_info "初始化数据库..."
    
    # 检查 mysql-init.sql 是否存在
    if [ ! -f "$SCRIPT_DIR/mysql-init.sql" ]; then
        log_error "找不到 mysql-init.sql 文件"
        exit 1
    fi
    
    # 执行初始化脚本
    if docker exec -i "$MYSQL_CONTAINER_NAME" mysql -uroot -p"$MYSQL_ROOT_PASSWORD" < "$SCRIPT_DIR/mysql-init.sql"; then
        log_success "数据库初始化完成"
    else
        log_error "数据库初始化失败"
        exit 1
    fi
}

# 安装 Go 依赖
install_dependencies() {
    log_info "安装 Go 依赖..."
    
    cd "$PROJECT_ROOT"
    
    if [ -f "go.mod" ]; then
        go mod download
        go mod tidy
        log_success "Go 依赖安装完成"
    else
        log_error "找不到 go.mod 文件"
        exit 1
    fi
}

# 编译后端代码
build_backend() {
    log_info "编译后端代码..."
    
    cd "$PROJECT_ROOT"
    
    # 创建 bin 目录
    mkdir -p bin
    
    # 编译
    if go build -o bin/vul-backend cmd/api/main.go; then
        log_success "编译完成: bin/vul-backend"
    else
        log_error "编译失败"
        exit 1
    fi
}

# 检查配置文件
check_config() {
    log_info "检查配置文件..."
    
    cd "$PROJECT_ROOT"
    
    if [ ! -f "configs/config.yaml" ]; then
        if [ -f "configs/config.example.yaml" ]; then
            log_warning "未找到 config.yaml，从 config.example.yaml 复制"
            cp configs/config.example.yaml configs/config.yaml
            log_success "配置文件已创建"
        else
            log_error "找不到配置文件"
            exit 1
        fi
    else
        log_success "配置文件检查通过"
    fi
}

# 启动后端服务
start_backend() {
    log_info "启动后端服务..."
    
    cd "$PROJECT_ROOT"
    
    # 检查 API 端口
    check_port "$API_PORT"
    
    # 启动服务
    nohup ./bin/vul-backend > logs/app.log 2>&1 &
    local backend_pid=$!
    
    log_info "后端服务 PID: $backend_pid"
    
    # 等待服务启动
    log_info "等待服务就绪..."
    sleep 3
    
    # 检查服务是否成功启动
    local max_attempts=10
    local attempt=0
    
    while [ $attempt -lt $max_attempts ]; do
        if curl -s http://localhost:$API_PORT/health &> /dev/null; then
            log_success "后端服务已启动"
            return 0
        fi
        
        # 检查进程是否还在运行
        if ! ps -p $backend_pid &> /dev/null; then
            log_error "后端服务启动失败，查看日志: logs/app.log"
            exit 1
        fi
        
        echo -n "."
        sleep 1
        ((attempt++))
    done
    
    log_warning "无法连接到健康检查端点，但服务可能正在启动"
}

# 停止后端服务
stop_backend() {
    log_info "停止后端服务..."
    
    local pids=$(pgrep -f "vul-backend" 2>/dev/null | grep -v $$ | grep -v grep || true)
    
    if [ -z "$pids" ]; then
        log_warning "未找到运行中的后端服务"
    else
        echo "$pids" | xargs kill -15 2>/dev/null || true
        sleep 2
        # 如果还在运行，强制终止
        pids=$(pgrep -f "vul-backend" 2>/dev/null | grep -v $$ | grep -v grep || true)
        if [ -n "$pids" ]; then
            echo "$pids" | xargs kill -9 2>/dev/null || true
        fi
        log_success "后端服务已停止"
    fi
}

# 停止 MySQL 容器
stop_mysql() {
    log_info "停止 MySQL 容器..."
    
    cd "$SCRIPT_DIR"
    
    if docker ps | grep -q "$MYSQL_CONTAINER_NAME"; then
        docker-compose down
        log_success "MySQL 容器已停止"
    else
        log_warning "MySQL 容器未运行"
    fi
}

# 打印服务信息
print_service_info() {
    echo ""
    echo -e "${GREEN}========================================${NC}"
    echo -e "${GREEN}   服务启动成功！${NC}"
    echo -e "${GREEN}========================================${NC}"
    echo ""
    echo -e "${BLUE}MySQL 数据库:${NC}"
    echo -e "  地址: localhost:$MYSQL_PORT"
    echo -e "  用户: root"
    echo -e "  密码: $MYSQL_ROOT_PASSWORD"
    echo -e "  数据库: $MYSQL_DATABASE"
    echo ""
    echo -e "${BLUE}后端 API:${NC}"
    echo -e "  地址: http://localhost:$API_PORT"
    echo -e "  健康检查: http://localhost:$API_PORT/health"
    echo -e "  版本信息: http://localhost:$API_PORT/version"
    echo -e "  API 文档: http://localhost:$API_PORT/api/v1/vulnerabilities"
    echo ""
    echo -e "${BLUE}日志文件:${NC}"
    echo -e "  后端日志: logs/app.log"
    echo ""
    echo -e "${BLUE}常用命令:${NC}"
    echo -e "  查看后端日志: tail -f logs/app.log"
    echo -e "  查看 MySQL 日志: docker logs -f $MYSQL_CONTAINER_NAME"
    echo -e "  停止服务: $0 stop"
    echo -e "  重启服务: $0 restart"
    echo -e "  查看状态: $0 status"
    echo ""
}

# 测试 API
test_api() {
    log_info "测试 API 连接..."
    
    echo ""
    echo -e "${BLUE}健康检查:${NC}"
    curl -s http://localhost:$API_PORT/health | python3 -m json.tool 2>/dev/null || curl -s http://localhost:$API_PORT/health
    
    echo ""
    echo -e "${BLUE}版本信息:${NC}"
    curl -s http://localhost:$API_PORT/version | python3 -m json.tool 2>/dev/null || curl -s http://localhost:$API_PORT/version
    echo ""
}

# 查看服务状态
check_status() {
    print_header "服务状态检查"
    
    # 检查 MySQL
    echo -e "${BLUE}MySQL 容器状态:${NC}"
    if docker ps | grep -q "$MYSQL_CONTAINER_NAME"; then
        echo -e "  ${GREEN}✓${NC} 运行中"
        docker ps --filter "name=$MYSQL_CONTAINER_NAME" --format "  容器 ID: {{.ID}}\n  状态: {{.Status}}\n  端口: {{.Ports}}"
    else
        echo -e "  ${RED}✗${NC} 未运行"
    fi
    echo ""
    
    # 检查后端
    echo -e "${BLUE}后端服务状态:${NC}"
    local pids=$(pgrep -f "vul-backend" 2>/dev/null | grep -v $$ | grep -v grep || true)
    if [ -n "$pids" ]; then
        echo -e "  ${GREEN}✓${NC} 运行中"
        echo "  进程 PID: $pids"
        
        # 检查端口
        if lsof -Pi :$API_PORT -sTCP:LISTEN -t &> /dev/null; then
            echo -e "  端口 $API_PORT: ${GREEN}监听中${NC}"
        fi
        
        # 测试健康检查
        if curl -s http://localhost:$API_PORT/health &> /dev/null; then
            echo -e "  健康检查: ${GREEN}正常${NC}"
        else
            echo -e "  健康检查: ${RED}失败${NC}"
        fi
    else
        echo -e "  ${RED}✗${NC} 未运行"
    fi
    echo ""
    
    # 检查端口占用
    echo -e "${BLUE}端口占用情况:${NC}"
    echo -e "  MySQL 端口 ($MYSQL_PORT):"
    if lsof -Pi :$MYSQL_PORT -sTCP:LISTEN -t &> /dev/null; then
        echo -e "    ${GREEN}✓${NC} 已占用"
    else
        echo -e "    ${RED}✗${NC} 空闲"
    fi
    
    echo -e "  API 端口 ($API_PORT):"
    if lsof -Pi :$API_PORT -sTCP:LISTEN -t &> /dev/null; then
        echo -e "    ${GREEN}✓${NC} 已占用"
    else
        echo -e "    ${RED}✗${NC} 空闲"
    fi
    echo ""
}

# 查看日志
view_logs() {
    print_header "查看服务日志"
    
    echo -e "${YELLOW}选择要查看的日志:${NC}"
    echo "  1) 后端应用日志"
    echo "  2) MySQL 容器日志"
    echo "  3) 同时查看所有日志"
    echo ""
    read -p "请选择 (1-3): " -n 1 -r
    echo ""
    
    case $REPLY in
        1)
            log_info "查看后端日志 (Ctrl+C 退出)"
            sleep 1
            tail -f "$PROJECT_ROOT/logs/app.log" 2>/dev/null || log_warning "日志文件不存在"
            ;;
        2)
            log_info "查看 MySQL 日志 (Ctrl+C 退出)"
            sleep 1
            docker logs -f "$MYSQL_CONTAINER_NAME" 2>/dev/null || log_warning "MySQL 容器未运行"
            ;;
        3)
            log_info "同时查看所有日志 (Ctrl+C 退出)"
            sleep 1
            # 使用 tmux 或 screen 同时查看，如果没有则顺序查看
            if command -v tmux &> /dev/null; then
                tmux new-session -d "tail -f $PROJECT_ROOT/logs/app.log" \; split-window -h "docker logs -f $MYSQL_CONTAINER_NAME" \; attach
            else
                tail -f "$PROJECT_ROOT/logs/app.log" 2>/dev/null &
                docker logs -f "$MYSQL_CONTAINER_NAME" 2>/dev/null
            fi
            ;;
        *)
            log_error "无效选择"
            ;;
    esac
}

# 清理所有数据（危险操作）
clean_all() {
    print_header "清理所有数据"
    
    echo -e "${RED}警告: 此操作将删除所有数据，包括:${NC}"
    echo "  - 停止并删除 MySQL 容器"
    echo "  - 删除 MySQL 数据卷"
    echo "  - 停止后端服务"
    echo "  - 清理编译文件"
    echo ""
    read -p "确认执行清理操作? (yes/no): " -r
    echo ""
    
    if [[ $REPLY == "yes" ]]; then
        stop_backend
        
        log_info "清理 Docker 容器和数据卷..."
        cd "$SCRIPT_DIR"
        docker-compose down -v 2>/dev/null || true
        log_success "Docker 资源已清理"
        
        log_info "清理编译文件..."
        cd "$PROJECT_ROOT"
        rm -rf bin/ 2>/dev/null || true
        log_success "编译文件已清理"
        
        log_success "所有数据已清理完成"
    else
        log_info "取消清理操作"
    fi
}

# 启动服务
start_service() {
    print_header "启动漏洞管理后端服务"
    
    # 创建日志目录
    mkdir -p "$PROJECT_ROOT/logs"
    
    # 1. 检查依赖
    check_dependencies
    
    # 2. 清理旧容器
    cleanup_old_containers
    
    # 3. 启动 MySQL
    start_mysql
    
    # 4. 初始化数据库
    init_database
    
    # 5. 检查配置文件
    check_config
    
    # 6. 安装依赖
    install_dependencies
    
    # 7. 编译代码
    build_backend
    
    # 8. 启动后端服务
    start_backend
    
    # 9. 打印服务信息
    print_service_info
    
    # 10. 测试 API
    test_api
    
    log_success "所有服务已成功启动！"
    echo ""
}

# 停止服务
stop_service() {
    print_header "停止漏洞管理后端服务"
    
    stop_backend
    stop_mysql
    
    echo ""
    log_success "所有服务已停止"
    echo ""
}

# 重启服务
restart_service() {
    print_header "重启漏洞管理后端服务"
    
    log_info "正在停止服务..."
    stop_backend
    # MySQL 保持运行，不重启
    
    log_info "正在重新编译..."
    cd "$PROJECT_ROOT"
    build_backend
    
    log_info "正在启动服务..."
    start_backend
    
    echo ""
    log_success "服务已重启！"
    echo ""
    
    # 测试 API
    test_api
}

# 主函数
main() {
    case "${1:-}" in
        start)
            start_service
            ;;
        stop)
            stop_service
            ;;
        restart)
            restart_service
            ;;
        status)
            check_status
            ;;
        logs)
            view_logs
            ;;
        clean)
            clean_all
            ;;
        *)
            print_usage
            exit 1
            ;;
    esac
}

# 捕获退出信号
trap 'log_warning "正在退出..."; exit 0' INT TERM

# 运行主函数
main "$@"