#!/bin/bash

# ============================================
# 数据库完整重建脚本
# 删除所有表并重新初始化（包含所有迁移）
# 
# 更新日志:
#   2026-01-23: 添加 all_tables.sql 整合文件，--simple 模式
#   2026-01-23: 添加 --drop-only, --init-only, --list-migrations 选项
#   2026-01-22: 添加 extract_tables 迁移支持
# ============================================

set -e

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
CYAN='\033[0;36m'
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"; }
log_step() { echo -e "${CYAN}[STEP]${NC} $1"; }

# 获取脚本所在目录
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
}

# 列出迁移脚本
list_migrations() {
    log_title "迁移脚本列表"
    
    if [ -d "$MIGRATIONS_DIR" ]; then
        local count=0
        echo ""
        printf "  %-5s %-60s %s\n" "序号" "文件名" "大小"
        echo "  -------------------------------------------------------------------------"
        for migration in $(ls -1 "${MIGRATIONS_DIR}"/*.sql 2>/dev/null | sort); do
            if [ -f "$migration" ]; then
                count=$((count + 1))
                local size=$(du -h "$migration" | cut -f1)
                printf "  %-5s %-60s %s\n" "$count" "$(basename $migration)" "$size"
            fi
        done
        echo ""
        log_info "共 ${count} 个迁移脚本"
    else
        log_warn "迁移目录不存在: ${MIGRATIONS_DIR}"
    fi
}

# 简单模式初始化（使用 all_tables.sql 单文件）
init_simple() {
    log_title "初始化数据库表（简单模式）"
    log_info "使用 all_tables.sql 单文件初始化所有表"
    execute_sql "${SQL_DIR}/all_tables.sql" "全部表结构 (21个表)"
}

# 初始化所有表（分文件模式）
init_all_tables() {
    log_title "初始化数据库表"
    
    local step=1
    
    # 1. 基础表（users, documents 等核心表）
    log_step "[$step/6] 基础表"
    execute_sql "${SQL_DIR}/init.sql" "基础表 (users, documents, elements, etc.)"
    step=$((step + 1))
    
    # 2. 图谱表（graph_nodes, graph_relations - 先于 supplement_tables）
    log_step "[$step/6] 图谱表"
    execute_sql "${SQL_DIR}/graph_tables.sql" "图谱表 (graph_nodes, graph_relations)"
    step=$((step + 1))
    
    # 3. 补充表（rules, data_sources, text_storage 等，不含 templates）
    log_step "[$step/6] 补充表"
    execute_sql "${SQL_DIR}/supplement_tables.sql" "补充表 (rules, data_sources, text_storage)"
    step=$((step + 1))
    
    # 4. RAG 表（text_chunks, vector_embeddings - 可能需要 pgvector）
    log_step "[$step/6] RAG 表"
    execute_sql "${SQL_DIR}/rag_tables_compatible.sql" "RAG 表 (text_chunks, vector_embeddings)"
    step=$((step + 1))
    
    # 5. 模板系统表（templates, source_files, variables, generations）
    log_step "[$step/6] 模板系统表"
    execute_sql "${SQL_DIR}/template_tables.sql" "模板系统表 (templates, source_files, variables, generations)"
    step=$((step + 1))
    
    # 6. 执行迁移文件（按文件名排序）
    log_step "[$step/6] 迁移脚本"
    log_title "执行迁移脚本"
    if [ -d "$MIGRATIONS_DIR" ]; then
        local migration_count=0
        local total_migrations=$(ls -1 "${MIGRATIONS_DIR}"/*.sql 2>/dev/null | wc -l)
        
        for migration in $(ls -1 "${MIGRATIONS_DIR}"/*.sql 2>/dev/null | sort); do
            if [ -f "$migration" ]; then
                migration_count=$((migration_count + 1))
                execute_sql "$migration" "[$migration_count/$total_migrations] $(basename $migration)"
            fi
        done
        log_info "共执行 ${migration_count} 个迁移脚本"
    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      不创建测试用户
    --drop-only         只删除表，不初始化
    --init-only         只初始化，不删除（用于全新数据库）
    --simple            使用 all_tables.sql 单文件初始化（推荐）
    --list-migrations   列出所有迁移脚本后退出

环境变量:
    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

    # 强制执行（跳过确认）
    ./rebuild_all.sh -f

    # 只删除所有表
    ./rebuild_all.sh --drop-only -f

    # 初始化全新数据库（不删除）
    ./rebuild_all.sh --init-only

    # 查看迁移脚本列表
    ./rebuild_all.sh --list-migrations

    # 简单模式（使用 all_tables.sql 单文件，推荐）
    ./rebuild_all.sh --simple -f

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

执行顺序:
    1. 删除所有表、序列、函数
    2. 执行 init.sql（用户、文档、元素等基础表）
    3. 执行 graph_tables.sql（图谱表）
    4. 执行 supplement_tables.sql（补充表，不含 templates）
    5. 执行 rag_tables_compatible.sql（RAG表）
    6. 执行 template_tables.sql（模板系统表 v2.0）
       - templates: 报告模板
       - source_files: 来源文件定义
       - variables: 模板变量
       - generations: 生成任务
    7. 执行 database/migrations/*.sql（迁移脚本，按名称排序）
    8. 创建测试用户
    9. 验证表创建

EOF
}

# 主函数
main() {
    CREATE_TEST_USER=true
    DROP_ONLY=false
    INIT_ONLY=false
    LIST_MIGRATIONS=false
    SIMPLE_MODE=false
    
    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
                ;;
            --drop-only)
                DROP_ONLY=true
                shift
                ;;
            --simple)
                SIMPLE_MODE=true
                shift
                ;;
            --init-only)
                INIT_ONLY=true
                shift
                ;;
            --list-migrations)
                LIST_MIGRATIONS=true
                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 ""
    
    # 列出迁移脚本模式
    if [ "$LIST_MIGRATIONS" = true ]; then
        list_migrations
        exit 0
    fi
    
    # 只删除模式
    if [ "$DROP_ONLY" = true ]; then
        drop_all_tables
        log_title "删除完成"
        log_info "所有表已删除，数据库已清空"
        exit 0
    fi
    
    # 只初始化模式
    if [ "$INIT_ONLY" = true ]; then
        if [ "$SIMPLE_MODE" = true ]; then
            init_simple
        else
            init_all_tables
        fi
        if [ "$CREATE_TEST_USER" = true ]; then
            create_test_user
        fi
        verify_tables
        log_title "初始化完成"
        log_info "所有表已创建并验证"
        exit 0
    fi
    
    # 完整重建模式
    drop_all_tables
    
    if [ "$SIMPLE_MODE" = true ]; then
        init_simple
    else
        init_all_tables
    fi
    
    if [ "$CREATE_TEST_USER" = true ]; then
        create_test_user
    fi
    
    verify_tables
    
    log_title "数据库重建完成"
    log_info "所有表已创建并验证"
}

main "$@"