#!/bin/bash
# quick_start.sh - 快速启动脚本
# 适配Docker部署的数据库环境

set -e

echo "========================================"
echo "AI+合规智能中枢 - 快速启动脚本"
echo "========================================"
echo ""

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

VENV_DIR=".venv"

# 检查Python版本
echo -e "${YELLOW}[1/8] 检查Python环境...${NC}"
if command -v python3 &> /dev/null; then
    PYTHON_CMD=python3
else
    PYTHON_CMD=python
fi

PYTHON_VERSION=$($PYTHON_CMD --version 2>&1 | awk '{print $2}')
echo "Python版本: $PYTHON_VERSION"

if [[ ! "$PYTHON_VERSION" =~ ^3\.1[0-9] ]]; then
    echo -e "${RED}需要Python 3.10+，当前版本: $PYTHON_VERSION${NC}"
    exit 1
fi
echo -e "${GREEN}Python环境检查通过${NC}"
echo ""

# 创建虚拟环境
echo -e "${YELLOW}[2/8] 创建虚拟环境...${NC}"
if [ -d "$VENV_DIR" ]; then
    echo "虚拟环境已存在: $VENV_DIR"
else
    echo "正在创建虚拟环境..."
    $PYTHON_CMD -m venv $VENV_DIR
    echo -e "${GREEN}虚拟环境创建成功${NC}"
fi

# 激活虚拟环境
source $VENV_DIR/bin/activate
echo "已激活虚拟环境: $VENV_DIR"
echo ""

# 安装依赖（使用国内镜像源）
echo -e "${YELLOW}[3/8] 安装Python依赖...${NC}"

# 配置pip使用清华镜像源（国内加速）
PIP_MIRROR="https://mirrors.aliyun.com/pypi/simple"

echo "使用镜像源: $PIP_MIRROR"
pip config set global.index-url $PIP_MIRROR -q

pip install --upgrade pip -q
pip install -r requirements.txt -q

if [ $? -eq 0 ]; then
    echo -e "${GREEN}依赖安装完成${NC}"
else
    echo -e "${RED}依赖安装失败，请检查requirements.txt${NC}"
    exit 1
fi
echo ""

# 检查Docker容器状态
echo -e "${YELLOW}[4/8] 检查Docker容器状态...${NC}"
REQUIRED_CONTAINERS="milvus minio redis postgres"

for container in $REQUIRED_CONTAINERS; do
    if docker ps --format '{{.Names}}' | grep -q "^${container}$"; then
        echo -e "${GREEN}✓ ${container} 容器运行正常${NC}"
    elif docker ps -a --format '{{.Names}}' | grep -q "^${container}$"; then
        echo -e "${YELLOW}⚠ ${container} 容器已停止，尝试启动...${NC}"
        docker start ${container}
        sleep 2
        if docker ps --format '{{.Names}}' | grep -q "^${container}$"; then
            echo -e "${GREEN}✓ ${container} 已启动${NC}"
        else
            echo -e "${RED}✗ ${container} 启动失败${NC}"
        fi
    else
        echo -e "${RED}✗ ${container} 容器不存在${NC}"
    fi
done
echo ""

# 检查PostgreSQL连接（通过Python）
echo -e "${YELLOW}[5/8] 检查PostgreSQL连接...${NC}"
python << 'EOF'
import sys
try:
    import psycopg2
    conn = psycopg2.connect(
        host="localhost",
        port=5432,
        user="postgresql",
        password="postgresql123456",
        database="postgres"
    )
    cur = conn.cursor()

    # 检查compliance_db是否存在
    cur.execute("SELECT 1 FROM pg_database WHERE datname='compliance_db'")
    exists = cur.fetchone()

    if not exists:
        cur.execute("CREATE DATABASE compliance_db")
        print("数据库 compliance_db 创建成功")
    else:
        print("数据库 compliance_db 已存在")

    conn.commit()
    cur.close()
    conn.close()
    print("PostgreSQL连接成功")
except Exception as e:
    print(f"PostgreSQL连接失败: {e}")
    sys.exit(1)
EOF

if [ $? -eq 0 ]; then
    echo -e "${GREEN}PostgreSQL服务运行正常${NC}"
else
    echo -e "${RED}PostgreSQL连接失败${NC}"
    exit 1
fi
echo ""

# 检查Redis连接（通过Python）
echo -e "${YELLOW}[6/8] 检查Redis连接...${NC}"
python << 'EOF'
import sys
try:
    import redis
    r = redis.Redis(
        host="localhost",
        port=6379,
        password="redis@123",
        decode_responses=True
    )
    result = r.ping()
    if result:
        print("Redis连接成功")
    else:
        sys.exit(1)
except Exception as e:
    print(f"Redis连接失败: {e}")
    sys.exit(1)
EOF

if [ $? -eq 0 ]; then
    echo -e "${GREEN}Redis服务运行正常${NC}"
else
    echo -e "${RED}Redis连接失败${NC}"
    exit 1
fi
echo ""

# 检查Milvus连接（通过Python）
echo -e "${YELLOW}[7/8] 检查Milvus连接...${NC}"
python << 'EOF'
import sys
try:
    from pymilvus import connections, utility
    connections.connect(host="localhost", port=19530)
    print("Milvus连接成功")
    connections.disconnect("default")
except Exception as e:
    print(f"Milvus连接失败: {e}")
    sys.exit(1)
EOF

if [ $? -eq 0 ]; then
    echo -e "${GREEN}Milvus服务运行正常${NC}"
else
    echo -e "${RED}Milvus连接失败${NC}"
    exit 1
fi
echo ""

# 检查MinIO连接（通过Python）
echo -e "${YELLOW}[8/8] 检查MinIO连接...${NC}"
python << 'EOF'
import sys
try:
    from minio import Minio
    client = Minio("localhost:9000", "minioadmin", "minioadmin", secure=False)

    # 检查bucket是否存在
    bucket = "compliance-docs"
    if not client.bucket_exists(bucket):
        client.make_bucket(bucket)
        print(f"MinIO bucket '{bucket}' 创建成功")
    else:
        print(f"MinIO bucket '{bucket}' 已存在")
    print("MinIO连接成功")
except Exception as e:
    print(f"MinIO连接失败: {e}")
    sys.exit(1)
EOF

if [ $? -eq 0 ]; then
    echo -e "${GREEN}MinIO服务运行正常${NC}"
else
    echo -e "${RED}MinIO连接失败${NC}"
    exit 1
fi
echo ""

# 预下载BGE-M3模型（可选）
echo -e "${YELLOW}[提示] BGE-M3嵌入模型...${NC}"
MODEL_CACHE=~/.cache/huggingface/hub/models--BAAI--bge-m3

if [ -d "$MODEL_CACHE" ]; then
    echo -e "${GREEN}BGE-M3模型已存在${NC}"
else
    echo -e "${YELLOW}模型未下载，首次运行时将自动下载（约2GB）${NC}"
    echo "手动预下载: python -c \"from FlagEmbedding import BGEM3FlagModel; BGEM3FlagModel('BAAI/bge-m3')\""
fi
echo ""

# 输出启动命令
echo "========================================"
echo -e "${GREEN}环境检查完成！${NC}"
echo "========================================"
echo ""
echo "虚拟环境: $VENV_DIR"
echo ""
echo "启动API服务:"
echo "  ./start_api.sh"
echo ""
echo "启动前端页面:"
echo "  ./start_frontend.sh"
echo "  访问: http://localhost:3000/prototype.html"
echo ""
echo "后台启动:"
echo "  ./start_api_background.sh"
echo ""
echo "API文档地址:"
echo "  http://localhost:8000/docs"
echo "  http://localhost:8000/health"
echo ""