#!/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}
DB_PASSWORD=${DB_PASSWORD:-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)"
PROJECT_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)"
SQL_DIR="${SCRIPT_DIR}"
MIGRATIONS_DIR="${PROJECT_ROOT}/database/migrations"

# 删除所有表
drop_all_tables() {
    log_title "删除所有表"
    
    if [ "$FORCE" != "yes" ]; then
        log_warn "这将删除数据库 ${DB_NAME} 中的所有表和数据！"
        read -p "是否继续？(yes/no): " confirm
        
        if [ "$confirm" != "yes" ]; then
            log_info "操作已取消"
            exit 0
        fi
    fi
    
    log_info "开始删除表..."
    
    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
\$\$;

-- 删除所有序列
DO \$\$
DECLARE
    r RECORD;
BEGIN
    FOR r IN (SELECT sequencename FROM pg_sequences WHERE schemaname = 'public')
    LOOP
        EXECUTE 'DROP SEQUENCE IF EXISTS public.' || quote_ident(r.sequencename) || ' CASCADE';
    END LOOP;
END
\$\$;

-- 删除所有函数
DO \$\$
DECLARE
    r RECORD;
BEGIN
    FOR r IN (SELECT proname, oidvectortypes(proargtypes) as args 
              FROM pg_proc 
              INNER JOIN pg_namespace ns ON (pg_proc.pronamespace = ns.oid)
              WHERE ns.nspname = 'public')
    LOOP
        EXECUTE 'DROP FUNCTION IF EXISTS public.' || quote_ident(r.proname) || '(' || r.args || ') CASCADE';
    END LOOP;
END
\$\$;

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

SELECT 'Tables remaining:' as status, COUNT(*) as count FROM pg_tables WHERE schemaname = 'public';
EOF
    
    log_info "所有表已删除"
}

# 执行 SQL 文件
execute_sql() {
    local file=$1
    local desc=$2
    
    if [ -f "$file" ]; then
        log_info "执行: ${desc} (${file})"
        PGPASSWORD=${DB_PASSWORD} psql -h ${DB_HOST} -p ${DB_PORT} -U ${DB_USER} -d ${DB_NAME} -f "$file" 2>&1 | grep -v "^NOTICE:" || true
        log_info "${desc} 完成"
    else
        log_warn "文件不存在: ${file}"
    fi
}

# 初始化所有表
init_all_tables() {
    log_title "初始化数据库表"
    
    # 1. 基础表（users, documents 等核心表）
    execute_sql "${SQL_DIR}/init.sql" "基础表 (users, documents, elements, etc.)"
    
    # 2. 图谱表（graph_nodes, graph_relations - 先于 supplement_tables）
    execute_sql "${SQL_DIR}/graph_tables.sql" "图谱表 (graph_nodes, graph_relations)"
    
    # 3. 补充表（rules, data_sources, templates, text_storage 等）
    execute_sql "${SQL_DIR}/supplement_tables.sql" "补充表 (rules, data_sources, templates, text_storage)"
    
    # 4. RAG 表（text_chunks, vector_embeddings - 可能需要 pgvector）
    execute_sql "${SQL_DIR}/rag_tables_compatible.sql" "RAG 表 (text_chunks, vector_embeddings)"
    
    # 5. 执行迁移文件（按文件名排序）
    log_title "执行迁移脚本"
    if [ -d "$MIGRATIONS_DIR" ]; then
        for migration in $(ls -1 "${MIGRATIONS_DIR}"/*.sql 2>/dev/null | sort); do
            if [ -f "$migration" ]; then
                execute_sql "$migration" "迁移: $(basename $migration)"
            fi
        done
    else
        log_warn "迁移目录不存在: ${MIGRATIONS_DIR}"
    fi
}

# 验证表创建
verify_tables() {
    log_title "验证表创建"
    
    PGPASSWORD=${DB_PASSWORD} psql -h ${DB_HOST} -p ${DB_PORT} -U ${DB_USER} -d ${DB_NAME} <<EOF
SELECT tablename as "表名"
FROM pg_tables 
WHERE schemaname = 'public'
ORDER BY tablename;

SELECT COUNT(*) as "表总数" FROM pg_tables WHERE schemaname = 'public';
EOF
}

# 创建测试用户
create_test_user() {
    log_title "创建测试用户"
    
    PGPASSWORD=${DB_PASSWORD} psql -h ${DB_HOST} -p ${DB_PORT} -U ${DB_USER} -d ${DB_NAME} <<EOF
-- 插入测试管理员用户（密码: Admin123!）
INSERT INTO users (id, username, email, password_hash, role, create_time, update_time)
VALUES (
    '00000000000000000000000000000001',
    'admin',
    'admin@lingyue.com',
    '\$2a\$10\$N.zmdr9k7uOCQb376NoUnuTJ8iUtkWBrq.VZbISbAi1L9sNPIiMMi',
    'admin',
    CURRENT_TIMESTAMP,
    CURRENT_TIMESTAMP
)
ON CONFLICT (id) DO NOTHING;

-- 插入测试普通用户（密码: User1234!）
INSERT INTO users (id, username, email, password_hash, role, create_time, update_time)
VALUES (
    '00000000000000000000000000000002',
    'testuser',
    'test@lingyue.com',
    '\$2a\$10\$5XvCJaLqDXJz.YqL8TnNKegHH7q6KM7YFzMVxLPNMvJx5f7mxnKQi',
    'user',
    CURRENT_TIMESTAMP,
    CURRENT_TIMESTAMP
)
ON CONFLICT (id) DO NOTHING;

SELECT id, username, email, role FROM users;
EOF
    
    log_info "测试用户创建完成"
    log_info "  admin / Admin123!"
    log_info "  testuser / User1234!"
}

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

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

选项:
    -h, --help      显示帮助
    -f, --force     跳过确认提示
    --no-test-user  不创建测试用户

环境变量:
    DB_HOST         数据库主机 (默认: localhost)
    DB_PORT         数据库端口 (默认: 5432)
    DB_NAME         数据库名称 (默认: lingyue_zhibao)
    DB_USER         数据库用户 (默认: postgres)
    DB_PASSWORD     数据库密码 (默认: postgres)

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

    # 指定数据库配置
    DB_PASSWORD=mypassword ./rebuild_all.sh

    # 强制执行（跳过确认）
    FORCE=yes ./rebuild_all.sh

    # 服务器上执行
    DB_USER=lingyue DB_PASSWORD=123123 ./rebuild_all.sh

执行顺序:
    1. 删除所有表、序列、函数
    2. 执行 init.sql（用户、文档、元素等基础表）
    3. 执行 supplement_tables.sql（补充表）
    4. 执行 graph_tables.sql（图谱表）
    5. 执行 rag_tables_compatible.sql（RAG表）
    6. 执行 database/migrations/*.sql（迁移脚本）
    7. 创建测试用户
    8. 验证表创建

EOF
}

# 主函数
main() {
    CREATE_TEST_USER=true
    
    while [[ $# -gt 0 ]]; do
        case $1 in
            -h|--help)
                show_help
                exit 0
                ;;
            -f|--force)
                FORCE=yes
                shift
                ;;
            --no-test-user)
                CREATE_TEST_USER=false
                shift
                ;;
            *)
                log_error "未知选项: $1"
                show_help
                exit 1
                ;;
        esac
    done
    
    log_title "数据库完整重建脚本"
    echo "数据库: ${DB_NAME}"
    echo "主机: ${DB_HOST}:${DB_PORT}"
    echo "用户: ${DB_USER}"
    echo "SQL目录: ${SQL_DIR}"
    echo "迁移目录: ${MIGRATIONS_DIR}"
    echo ""
    
    drop_all_tables
    init_all_tables
    
    if [ "$CREATE_TEST_USER" = true ]; then
        create_test_user
    fi
    
    verify_tables
    
    log_title "数据库重建完成"
    log_info "所有表已创建并验证"
}

main "$@"