#!/bin/bash
# 数据库完整重建脚本：删除所有表后执行 database/init.sql（单文件初始化）

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)"
INIT_SQL="${PROJECT_ROOT}/database/init.sql"

execute_sql() {
    local file="$1"
    local desc="${2:-$(basename "$file")}"
    if [ ! -f "$file" ]; then
        log_error "文件不存在: $file"
        return 1
    fi
    log_step "执行: $desc"
    PGPASSWORD=${DB_PASSWORD} psql -h ${DB_HOST} -p ${DB_PORT} -U ${DB_USER} -d ${DB_NAME} -f "$file" || return 1
}

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';
EOF
    log_info "删除完成"
}

init_database() {
    log_title "初始化数据库"
    execute_sql "${INIT_SQL}" "database/init.sql（全表结构）"
}

create_test_user() {
    log_title "创建测试用户"
    PGPASSWORD=${DB_PASSWORD} psql -h ${DB_HOST} -p ${DB_PORT} -U ${DB_USER} -d ${DB_NAME} <<'EOF'
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;
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 "测试用户: admin/Admin123!  testuser/User1234!"
}

show_help() {
    cat <<EOF
用法: ./rebuild_all.sh [选项]
选项:
  -h, --help       显示帮助
  -f, --force      跳过确认
  --drop-only      只删除表
  --init-only      只执行 database/init.sql（不删除）
  --no-test-user   不创建测试用户

环境变量: DB_HOST, DB_PORT, DB_NAME, DB_USER, DB_PASSWORD
EOF
}

CREATE_TEST_USER=true
DROP_ONLY=false
INIT_ONLY=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 ;;
        --init-only) INIT_ONLY=true; shift ;;
        *) log_error "未知选项: $1"; show_help; exit 1 ;;
    esac
done

log_title "数据库重建"
echo "数据库: ${DB_NAME} @ ${DB_HOST}:${DB_PORT}"
echo "初始化脚本: ${INIT_SQL}"
echo ""

if [ "$DROP_ONLY" = true ]; then
    drop_all_tables
    exit 0
fi

if [ "$INIT_ONLY" = true ]; then
    init_database
    [ "$CREATE_TEST_USER" = true ] && create_test_user
    log_info "初始化完成"
    exit 0
fi

drop_all_tables
init_database
[ "$CREATE_TEST_USER" = true ] && create_test_user
log_info "重建完成"