#!/bin/bash

# ============================================
# 仅初始化补充表（不删除现有表）
# ============================================

set -e

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

# 数据库配置
DB_HOST=${DB_HOST:-localhost}
DB_PORT=${DB_PORT:-5432}
DB_NAME=${DB_NAME:-lingyue_zhibao}
DB_USER=${DB_USER:-postgres}

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_title() { echo -e "\n${BLUE}========== $1 ==========${NC}\n"; }

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
SQL_DIR="${SCRIPT_DIR}"

log_title "初始化补充表"
echo "数据库: ${DB_NAME}"
echo "主机: ${DB_HOST}:${DB_PORT}"
echo "用户: ${DB_USER}"
echo ""

log_info "执行 supplement_tables.sql..."
PGPASSWORD=${DB_PASSWORD} psql -h ${DB_HOST} -p ${DB_PORT} -U ${DB_USER} -d ${DB_NAME} -f "${SQL_DIR}/supplement_tables.sql"

if [ $? -eq 0 ]; then
    log_info "补充表初始化成功"
    
    log_info "验证表创建..."
    PGPASSWORD=${DB_PASSWORD} psql -h ${DB_HOST} -p ${DB_PORT} -U ${DB_USER} -d ${DB_NAME} <<EOF
SELECT tablename 
FROM pg_tables 
WHERE schemaname = 'public' 
AND tablename IN ('text_storage', 'vector_embeddings', 'knowledge_base', 'rag_sessions', 'rag_messages')
ORDER BY tablename;
EOF
    
    log_title "完成"
else
    log_error "补充表初始化失败"
    exit 1
fi