#!/bin/bash

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

# 设置脚本所在目录为工作目录
SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
SQL_DIR="${SCRIPT_DIR}/sql"

# 打印带时间戳的日志
log() {
    local level=$1
    local message=$2
    local color=$NC
    
    case $level in
        "INFO") color=$BLUE;;
        "SUCCESS") color=$GREEN;;
        "WARN") color=$YELLOW;;
        "ERROR") color=$RED;;
    esac
    
    echo -e "[$(date '+%Y-%m-%d %H:%M:%S')] ${color}${level}${NC}: ${message}"
}

# 数据库连接信息（可通过环境变量覆盖）
DB_HOST=${POSTGRES_HOST:-"localhost"}
DB_PORT=${POSTGRES_PORT:-"5432"}
DB_NAME=${POSTGRES_DB:-"urban-lifeline"}
DB_USER=${POSTGRES_USER:-"postgres"}
DB_PASSWORD=${POSTGRES_PASSWORD:-"postgres"}

# 设置 PSQL 环境变量以支持中文
export PGCLIENTENCODING=UTF8

# 检查psql命令是否可用
check_psql() {
    if ! command -v psql &> /dev/null; then
        echo -e "${RED}Error: psql command not found. Please install PostgreSQL client.${NC}"
        exit 1
    fi
}

# 检查并创建数据库用户
check_and_create_user() {
    local new_user=$1
    local new_password=$2
    
    # 使用 postgres 用户执行
    if sudo -u postgres psql -c "SELECT 1 FROM pg_roles WHERE rolname = '$new_user'" | grep -q 1; then
        echo -e "${GREEN}User $new_user already exists${NC}"
    else
        echo -e "${YELLOW}Creating user $new_user...${NC}"
        sudo -u postgres psql -c "CREATE USER $new_user WITH PASSWORD '$new_password' CREATEDB;"
        if [ $? -eq 0 ]; then
            echo -e "${GREEN}User $new_user created successfully${NC}"
        else
            echo -e "${RED}Failed to create user $new_user${NC}"
            exit 1
        fi
    fi
}

# 检查数据库连接
check_db_connection() {
    # 首先尝试以当前用户身份连接
    if ! psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "postgres" -c '\q' &> /dev/null; then
        echo -e "${YELLOW}Could not connect with current settings, attempting to create user...${NC}"
        # 创建用户并设置权限
        check_and_create_user "$DB_USER" "$DB_PASSWORD"
        
        # 再次检查连接
        if ! psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "postgres" -c '\q' &> /dev/null; then
            echo -e "${RED}Error: Could not connect to PostgreSQL server.${NC}"
            echo "Please check your connection settings:"
            echo "Host: $DB_HOST"
            echo "Port: $DB_PORT"
            echo "User: $DB_USER"
            exit 1
        fi
    fi
}

# 执行SQL文件
execute_sql_file() {
    local sql_file=$1
    if [ ! -f "$sql_file" ]; then
        echo -e "${RED}Error: SQL file not found: $sql_file${NC}"
        return 1
    fi
    
    echo -e "${YELLOW}Executing SQL file: $sql_file${NC}"
    PGPASSWORD=$DB_PASSWORD psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAME" -f "$sql_file"
    local status=$?
    if [ $status -eq 0 ]; then
        echo -e "${GREEN}Successfully executed: $sql_file${NC}"
    else
        echo -e "${RED}Failed to execute: $sql_file${NC}"
        return $status
    fi
}

# 初始化数据库
init() {
    echo -e "${YELLOW}Initializing database...${NC}"
    
    # 执行完整的初始化脚本
    log "INFO" "Executing initialization script..."
    # Run from inside the SQL_DIR so relative \i includes in initAll.sql (like createDB.sql)
    # resolve relative to the SQL directory.
    (
        if [ ! -d "$SQL_DIR" ]; then
            echo -e "${RED}Error: SQL directory not found: $SQL_DIR${NC}"
            exit 1
        fi
        cd "$SQL_DIR" || exit 1
        PGPASSWORD=$DB_PASSWORD psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "postgres" -v ON_ERROR_STOP=1 -f "initAll.sql"
    )
    
    if [ $? -eq 0 ]; then
        log "SUCCESS" "Database initialization completed successfully"
    else
        log "ERROR" "Database initialization failed"
        return 1
    fi
    
    if [ $? -ne 0 ]; then
        echo -e "${RED}Failed to create database.${NC}"
        return 1
    fi
    
    # 2. 创建扩展和设置搜索路径
    echo -e "${YELLOW}Creating extensions...${NC}"
    check_extensions_availability() {
        # 检查服务器上是否存在需创建的扩展
        local missing=()
        local exts=("uuid-ossp" "pg_trgm" "btree_gist")
        for ext in "${exts[@]}"; do
            # 查询 pg_available_extensions 来判断扩展是否已安装到服务器目录
            if ! PGPASSWORD=$DB_PASSWORD psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAME" -tAc "SELECT 1 FROM pg_available_extensions WHERE name = '$ext';" | grep -q 1; then
                missing+=("$ext")
            fi
        done

        if [ ${#missing[@]} -ne 0 ]; then
            echo -e "${RED}Error: The following server-side extensions are not available: ${missing[*]}${NC}"
            echo "If you compiled PostgreSQL from source, you need to build and install the contrib modules into the server's installation prefix. Example steps:"
            echo "  # 在 PostgreSQL 源码目录下运行："
            echo "  cd /path/to/postgresql-source/contrib"
            echo "  make"
            echo "  sudo make install"
            echo "或者只安装缺失的模块（例如 uuid-ossp）："
            echo "  cd /path/to/postgresql-source/contrib/uuid-ossp"
            echo "  make"
            echo "  sudo make install"
            echo "安装完成后，重启 PostgreSQL 服务并重新运行此脚本："
            echo "  sudo systemctl restart postgresql"
            echo "如果你使用的是容器或自定义路径，请确保将编译安装的扩展安装到 PostgreSQL 的 \$(pg_config --sharedir)/extension 目录下。"
            return 1
        fi
        return 0
    }

    # 检查扩展可用性，若缺失则给出建议并退出
    if ! check_extensions_availability; then
        return 1
    fi

    PGPASSWORD=$DB_PASSWORD psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAME" -c "
            CREATE EXTENSION IF NOT EXISTS \"uuid-ossp\";
            CREATE EXTENSION IF NOT EXISTS \"pg_trgm\";
            CREATE EXTENSION IF NOT EXISTS \"btree_gist\";"
    
    if [ $? -ne 0 ]; then
        echo -e "${RED}Failed to create extensions.${NC}"
        return 1
    fi
    
    # 3. 逐个执行初始化SQL文件
    echo -e "${YELLOW}Initializing tables...${NC}"
    while IFS= read -r line || [[ -n "$line" ]]; do
        # 跳过注释和空行
        [[ $line =~ ^--.*$ ]] && continue
        [[ -z "${line// }" ]] && continue
        
        # 从 \i 命令中提取文件名
        if [[ $line =~ \\i[[:space:]]+([^[:space:]]+) ]]; then
            sql_file="${SQL_DIR}/${BASH_REMATCH[1]}"
            if [[ $sql_file != *"createDB.sql"* ]]; then  # 跳过createDB.sql
                echo -e "${YELLOW}Executing: $sql_file${NC}"
                PGPASSWORD=$DB_PASSWORD psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAME" -f "$sql_file"
                if [ $? -ne 0 ]; then
                    echo -e "${RED}Failed to execute: $sql_file${NC}"
                    return 1
                fi
            fi
        fi
    done < "${SQL_DIR}/initAll.sql"
    
    # 4. 设置搜索路径
    echo -e "${YELLOW}Setting search path...${NC}"
    PGPASSWORD=$DB_PASSWORD psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAME" -c "
        ALTER DATABASE $DB_NAME SET search_path TO sys, public;"
    
    echo -e "${GREEN}Database initialization completed successfully.${NC}"
    return 0
}

# 重新初始化数据库
reinit() {
    echo -e "${YELLOW}Reinitializing database...${NC}"
    delete
    init
}

# 删除数据库
delete() {
    echo -e "${YELLOW}Deleting database...${NC}"
    
    # 确保没有活动连接
    PGPASSWORD=$DB_PASSWORD psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "postgres" -c "
        SELECT pg_terminate_backend(pg_stat_activity.pid)
        FROM pg_stat_activity
        WHERE pg_stat_activity.datname = '$DB_NAME'
        AND pid <> pg_backend_pid();"
    
    # 删除数据库
    PGPASSWORD=$DB_PASSWORD psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "postgres" -c "DROP DATABASE IF EXISTS $DB_NAME;"
    
    echo -e "${GREEN}Database deleted.${NC}"
}

# 显示帮助信息
show_help() {
    echo "Usage: $0 {init|reinit|delete}"
    echo "Commands:"
    echo "  init    Initialize the database"
    echo "  reinit  Reinitialize the database (delete and create)"
    echo "  delete  Delete the database"
}

# 主函数
main() {
    check_psql
    check_db_connection
    
    case "$1" in
        "init")
            init
            ;;
        "reinit")
            reinit
            ;;
        "delete")
            delete
            ;;
        *)
            show_help
            exit 1
            ;;
    esac
}

# Call main with all passed arguments so the script runs when invoked
main "$@"