#!/bin/bash

echo "======================================"
echo "  Chroma 向量数据库部署脚本"
echo "======================================"
echo ""

DEPLOY_DIR="/opt/chroma"
DATA_DIR="/data/chroma"

echo "[1/5] 检查环境..."
if ! command -v docker &> /dev/null; then
    echo "错误: Docker 未安装，请先安装 Docker"
    exit 1
fi

if ! command -v docker-compose &> /dev/null; then
    echo "错误: Docker Compose 未安装，请先安装 Docker Compose"
    exit 1
fi

echo "[2/5] 创建目录..."
mkdir -p $DEPLOY_DIR
mkdir -p $DATA_DIR/chroma_data

echo "[3/5] 生成配置文件..."

cat > $DEPLOY_DIR/docker-compose.yml << 'EOF'
version: '3.8'

services:
  chroma:
    image: chromadb/chroma:latest
    container_name: chroma_server
    ports:
      - "8000:8000"
    volumes:
      - /data/chroma/chroma_data:/chroma/chroma_data
    environment:
      - CHROMA_SERVER_HOST=0.0.0.0
      - CHROMA_SERVER_HTTP_PORT=8000
      - CHROMA_SERVER_AUTH_CREDENTIALS=${CHROMA_AUTH_TOKEN}
      - CHROMA_SERVER_AUTH_PROVIDER=chromadb.auth.token.TokenAuthProvider
      - ANONYMIZED_TELEMETRY=False
    restart: unless-stopped
    healthcheck:
      test: ["CMD", "curl", "-f", "http://localhost:8000/api/v1/heartbeat"]
      interval: 30s
      timeout: 10s
      retries: 3
    networks:
      - chroma_network

networks:
  chroma_network:
    driver: bridge
EOF

AUTH_TOKEN=$(openssl rand -hex 32)

cat > $DEPLOY_DIR/.env << EOF
CHROMA_AUTH_TOKEN=$AUTH_TOKEN
EMBEDDING_API_KEY=your_embedding_api_key_here
OPENAI_API_KEY=your_openai_api_key_here
EOF

echo "[4/5] 启动服务..."
cd $DEPLOY_DIR
docker-compose up -d

echo "[5/5] 等待服务启动..."
sleep 10

echo ""
echo "======================================"
echo "  部署完成！"
echo "======================================"
echo ""
echo "服务地址: http://localhost:8000"
echo "认证Token: $AUTH_TOKEN"
echo ""
echo "验证命令: curl http://localhost:8000/api/v1/heartbeat"
echo ""
echo "配置文件位置: $DEPLOY_DIR/.env"
echo ""