#!/bin/bash

# RAG 系统 Docker 启动脚本
# 用法: ./start.sh [dev|prod]

set -e

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

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

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

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

# 检查 Docker 是否安装
check_docker() {
    if ! command -v docker &> /dev/null; then
        print_error "Docker 未安装，请先安装 Docker"
        exit 1
    fi
    
    # 检查 docker-compose 或 docker compose
    if command -v docker-compose &> /dev/null; then
        DOCKER_COMPOSE_CMD="docker-compose"
    elif docker compose version &> /dev/null; then
        DOCKER_COMPOSE_CMD="docker compose"
    else
        print_error "Docker Compose 未安装，请先安装 Docker Compose"
        exit 1
    fi
    
    print_info "Docker 和 Docker Compose 已安装"
}

# 检查 .env 文件
check_env_file() {
    if [ ! -f "../.env" ]; then
        print_warn ".env 文件不存在，从 .env.example 复制"
        if [ -f "../.env.example" ]; then
            cp ../.env.example ../.env
            print_info ".env 文件已创建，请根据需要修改配置"
        else
            print_error ".env.example 文件不存在"
            exit 1
        fi
    fi
}

# 启动开发环境
start_dev() {
    print_info "启动开发环境..."
    $DOCKER_COMPOSE_CMD -f docker-compose.dev.yml up -d --build
    
    print_info "等待服务启动..."
    sleep 10
    
    print_info "开发环境已启动！"
    echo ""
    echo "访问地址："
    echo "  - API 文档: http://localhost:8000/docs"
    echo "  - pgAdmin: http://localhost:5050 (admin@example.com / admin)"
    echo ""
    echo "查看日志："
    echo "  $DOCKER_COMPOSE_CMD -f docker-compose.dev.yml logs -f"
    echo ""
    echo "停止服务："
    echo "  $DOCKER_COMPOSE_CMD -f docker-compose.dev.yml down"
}

# 启动生产环境
start_prod() {
    print_info "启动生产环境..."
    $DOCKER_COMPOSE_CMD -f docker-compose.yml up -d --build
    
    print_info "等待服务启动..."
    sleep 10
    
    print_info "生产环境已启动！"
    echo ""
    echo "访问地址："
    echo "  - API 文档: http://localhost:8000/docs"
    echo ""
    echo "查看日志："
    echo "  $DOCKER_COMPOSE_CMD -f docker-compose.yml logs -f"
    echo ""
    echo "停止服务："
    echo "  $DOCKER_COMPOSE_CMD -f docker-compose.yml down"
}

# 停止服务
stop_services() {
    local env=$1
    
    if [ "$env" = "dev" ]; then
        print_info "停止开发环境..."
        $DOCKER_COMPOSE_CMD -f docker-compose.dev.yml down
    elif [ "$env" = "prod" ]; then
        print_info "停止生产环境..."
        $DOCKER_COMPOSE_CMD -f docker-compose.yml down
    else
        print_info "停止所有服务..."
        $DOCKER_COMPOSE_CMD -f docker-compose.dev.yml down 2>/dev/null || true
        $DOCKER_COMPOSE_CMD -f docker-compose.yml down 2>/dev/null || true
    fi
    
    print_info "服务已停止"
}

# 显示帮助信息
show_help() {
    echo "用法: $0 [命令]"
    echo ""
    echo "命令："
    echo "  dev       启动开发环境"
    echo "  prod      启动生产环境"
    echo "  stop      停止所有服务"
    echo "  stop-dev  停止开发环境"
    echo "  stop-prod 停止生产环境"
    echo "  help      显示帮助信息"
    echo ""
    echo "示例："
    echo "  $0 dev        # 启动开发环境"
    echo "  $0 prod       # 启动生产环境"
    echo "  $0 stop       # 停止所有服务"
}

# 主函数
main() {
    # 切换到 docker 目录
    cd "$(dirname "$0")"
    
    local command=${1:-dev}
    
    case $command in
        dev)
            check_docker
            check_env_file
            start_dev
            ;;
        prod)
            check_docker
            check_env_file
            start_prod
            ;;
        stop)
            check_docker
            stop_services
            ;;
        stop-dev)
            check_docker
            stop_services "dev"
            ;;
        stop-prod)
            check_docker
            stop_services "prod"
            ;;
        help|--help|-h)
            show_help
            ;;
        *)
            print_error "未知命令: $command"
            show_help
            exit 1
            ;;
    esac
}

# 运行主函数
main "$@"