#!/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}"

# 检查 SQL 文件是否存在
check_sql_files() {
    log_title "检查 SQL 文件"
    
    local missing=false
    
    for file in init.sql supplement_tables.sql rag_tables.sql; do
        if [ -f "${SQL_DIR}/${file}" ]; then
            log_info "找到: ${file}"
        else
            log_error "缺失: ${file}"
            missing=true
        fi
    done
    
    if [ "$missing" = true ]; then
        log_error "请确保在 backend/sql 目录下运行此脚本"
        exit 1
    fi
}

# 删除所有表
drop_all_tables() {
    log_title "删除所有表"
    
    log_warn "这将删除数据库中的所有表和数据！"
    read -p "是否继续？(yes/no): " confirm
    
    if [ "$confirm" != "yes" ]; then
        log_info "操作已取消"
        exit 0
    fi
    
    log_info "开始删除表..."
    
    # 删除所有表的 SQL
    PGPASSWORD=${DB_PASSWORD} psql -h ${DB_HOST} -p ${DB_PORT} -U ${DB_USER} -d ${DB_NAME} <<EOF
-- 禁用外键约束检查
SET session_replication_role = 'replica';

-- 删除所有表
DO \$\$
DECLARE
    r RECORD;
BEGIN
    FOR r IN (SELECT tablename FROM pg_tables WHERE schemaname = 'public')
    LOOP
        EXECUTE 'DROP TABLE IF EXISTS public.' || quote_ident(r.tablename) || ' CASCADE';
        RAISE NOTICE 'Dropped table: %', r.tablename;
    END LOOP;
END
\$\$;

-- 恢复外键约束检查
SET session_replication_role = 'origin';

-- 显示剩余的表
SELECT tablename FROM pg_tables WHERE schemaname = 'public';
EOF
    
    log_info "所有表已删除"
}

# 初始化基础表
init_base_tables() {
    log_title "初始化基础表 (init.sql)"
    
    PGPASSWORD=${DB_PASSWORD} psql -h ${DB_HOST} -p ${DB_PORT} -U ${DB_USER} -d ${DB_NAME} -f "${SQL_DIR}/init.sql"
    
    if [ $? -eq 0 ]; then
        log_info "基础表初始化成功"
    else
        log_error "基础表初始化失败"
        exit 1
    fi
}

# 初始化补充表
init_supplement_tables() {
    log_title "初始化补充表 (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 "补充表初始化成功"
    else
        log_error "补充表初始化失败"
        exit 1
    fi
}

# 初始化 RAG 表
init_rag_tables() {
    log_title "初始化 RAG 表 (rag_tables.sql)"
    
    PGPASSWORD=${DB_PASSWORD} psql -h ${DB_HOST} -p ${DB_PORT} -U ${DB_USER} -d ${DB_NAME} -f "${SQL_DIR}/rag_tables.sql" 2>&1
    
    if [ $? -eq 0 ]; then
        log_info "RAG 表初始化成功"
    else
        log_warn "RAG 表初始化可能失败（如果未安装 pgvector 扩展是正常的）"
    fi
}

# 验证表创建
verify_tables() {
    log_title "验证表创建"
    
    log_info "查询所有表..."
    PGPASSWORD=${DB_PASSWORD} psql -h ${DB_HOST} -p ${DB_PORT} -U ${DB_USER} -d ${DB_NAME} <<EOF
SELECT 
    schemaname,
    tablename,
    pg_size_pretty(pg_total_relation_size(schemaname||'.'||tablename)) AS size
FROM pg_tables 
WHERE schemaname = 'public'
ORDER BY tablename;
EOF
    
    log_info "表统计..."
    PGPASSWORD=${DB_PASSWORD} psql -h ${DB_HOST} -p ${DB_PORT} -U ${DB_USER} -d ${DB_NAME} <<EOF
SELECT COUNT(*) as table_count FROM pg_tables WHERE schemaname = 'public';
EOF
}

# 显示帮助
show_help() {
    cat <<EOF
数据库重建脚本

用法: ./rebuild_database.sh [选项]

环境变量:
    DB_HOST         数据库主机 (默认: localhost)
    DB_PORT         数据库端口 (默认: 5432)
    DB_NAME         数据库名称 (默认: lingyue_zhibao)
    DB_USER         数据库用户 (默认: postgres)
    DB_PASSWORD     数据库密码 (可选，如果未设置将提示输入)

示例:
    # 使用默认配置
    ./rebuild_database.sh

    # 指定数据库配置
    DB_HOST=localhost DB_USER=postgres DB_PASSWORD=123456 ./rebuild_database.sh

    # 仅初始化新表（不删除）
    ./init_supplement_only.sh

注意:
    - 此脚本会删除所有现有表和数据
    - 确保已备份重要数据
    - 需要数据库管理员权限

执行步骤:
    1. 删除所有表
    2. 执行 init.sql（基础表）
    3. 执行 supplement_tables.sql（补充表）
    4. 执行 rag_tables.sql（RAG表）
    5. 验证表创建

EOF
}

# 主函数
main() {
    if [ "$1" = "-h" ] || [ "$1" = "--help" ]; then
        show_help
        exit 0
    fi
    
    log_title "数据库重建脚本"
    echo "数据库: ${DB_NAME}"
    echo "主机: ${DB_HOST}:${DB_PORT}"
    echo "用户: ${DB_USER}"
    echo ""
    
    check_sql_files
    drop_all_tables
    init_base_tables
    init_supplement_tables
    init_rag_tables
    verify_tables
    
    log_title "数据库重建完成"
    log_info "所有表已创建并验证"
}

main "$@"