#!/bin/bash

# H3C ONU设备管理系统 - 数据库初始化脚本
# 用于初始化数据库和创建基础数据

set -e

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

# 加载环境变量
if [ -f "../.env" ]; then
    set -a
    source ../.env
    set +a
else
    echo -e "${RED}错误: .env 文件不存在${NC}"
    echo "请先复制 .env.example 为 .env 并配置数据库连接"
    exit 1
fi

# 检查数据库连接字符串
if [ -z "$DATABASE_URL" ]; then
    echo -e "${RED}错误: DATABASE_URL 未设置${NC}"
    exit 1
fi

# 从DATABASE_URL解析连接信息
parse_db_url() {
    local url="$1"
    # 格式: postgresql://username:password@host:port/database
    local pattern="postgresql://([^:]+):([^@]+)@([^:]+):([^/]+)/(.+)"
    
    if [[ $url =~ $pattern ]]; then
        DB_USER="${BASH_REMATCH[1]}"
        DB_PASSWORD="${BASH_REMATCH[2]}"
        DB_HOST="${BASH_REMATCH[3]}"
        DB_PORT="${BASH_REMATCH[4]}"
        DB_NAME="${BASH_REMATCH[5]}"
    else
        echo -e "${RED}错误: 无法解析DATABASE_URL${NC}"
        exit 1
    fi
}

parse_db_url "$DATABASE_URL"

# 函数：检查PostgreSQL连接
check_postgres_connection() {
    echo -e "${BLUE}检查PostgreSQL连接...${NC}"
    
    if PGPASSWORD="$DB_PASSWORD" psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d postgres -c "SELECT 1;" > /dev/null 2>&1; then
        echo -e "${GREEN}✓ PostgreSQL连接成功${NC}"
        return 0
    else
        echo -e "${RED}✗ PostgreSQL连接失败${NC}"
        return 1
    fi
}

# 函数：检查数据库是否存在
check_database_exists() {
    echo -e "${BLUE}检查数据库 '$DB_NAME' 是否存在...${NC}"
    
    if PGPASSWORD="$DB_PASSWORD" psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAME" -c "SELECT 1;" > /dev/null 2>&1; then
        echo -e "${GREEN}✓ 数据库已存在${NC}"
        return 0
    else
        echo -e "${YELLOW}数据库不存在${NC}"
        return 1
    fi
}

# 函数：创建数据库
create_database() {
    echo -e "${BLUE}创建数据库 '$DB_NAME'...${NC}"
    
    # 创建数据库
    PGPASSWORD="$DB_PASSWORD" psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d postgres <<EOF
CREATE DATABASE $DB_NAME;
EOF
    
    if [ $? -eq 0 ]; then
        echo -e "${GREEN}✓ 数据库创建成功${NC}"
    else
        echo -e "${RED}✗ 数据库创建失败${NC}"
        exit 1
    fi
}

# 函数：创建数据库用户（如果不存在）
create_database_user() {
    echo -e "${BLUE}检查数据库用户 '$DB_USER'...${NC}"
    
    # 检查用户是否存在
    USER_EXISTS=$(PGPASSWORD="$DB_PASSWORD" psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d postgres -tAc "SELECT 1 FROM pg_roles WHERE rolname='$DB_USER'")
    
    if [ "$USER_EXISTS" = "1" ]; then
        echo -e "${GREEN}✓ 用户已存在${NC}"
    else
        echo -e "${BLUE}创建用户 '$DB_USER'...${NC}"
        
        # 需要以postgres用户创建新用户
        # 这里假设有postgres用户的密码或者使用trust认证
        echo -e "${YELLOW}注意: 需要以postgres用户权限创建新用户${NC}"
        echo "请手动执行以下命令:"
        echo "  psql -h $DB_HOST -p $DB_PORT -U postgres -d postgres"
        echo "  CREATE USER $DB_USER WITH PASSWORD '$DB_PASSWORD';"
        echo "  GRANT ALL PRIVILEGES ON DATABASE $DB_NAME TO $DB_USER;"
        
        read -p "是否已手动创建用户? (y/n): " -n 1 -r
        echo
        if [[ ! $REPLY =~ ^[Yy]$ ]]; then
            echo -e "${YELLOW}跳过用户创建${NC}"
        fi
    fi
}

# 函数：授予权限
grant_privileges() {
    echo -e "${BLUE}授予用户权限...${NC}"
    
    # 授予数据库权限
    PGPASSWORD="$DB_PASSWORD" psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d postgres <<EOF
GRANT ALL PRIVILEGES ON DATABASE $DB_NAME TO $DB_USER;
EOF
    
    # 授予schema权限
    PGPASSWORD="$DB_PASSWORD" psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAME" <<EOF
GRANT ALL ON SCHEMA public TO $DB_USER;
GRANT ALL PRIVILEGES ON ALL TABLES IN SCHEMA public TO $DB_USER;
GRANT ALL PRIVILEGES ON ALL SEQUENCES IN SCHEMA public TO $DB_USER;
GRANT ALL PRIVILEGES ON ALL FUNCTIONS IN SCHEMA public TO $DB_USER;
ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT ALL ON TABLES TO $DB_USER;
ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT ALL ON SEQUENCES TO $DB_USER;
ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT ALL ON FUNCTIONS TO $DB_USER;
EOF
    
    echo -e "${GREEN}✓ 权限授予成功${NC}"
}

# 函数：运行数据库迁移
run_migrations() {
    echo -e "${BLUE}运行数据库迁移...${NC}"
    
    # 检查后端目录是否存在
    if [ ! -d "../../backend" ]; then
        echo -e "${YELLOW}警告: 后端目录不存在，跳过迁移${NC}"
        return
    fi
    
    # 进入后端目录
    cd "../../backend"
    
    # 检查Alembic配置
    if [ ! -f "alembic.ini" ]; then
        echo -e "${YELLOW}警告: alembic.ini 不存在，跳过迁移${NC}"
        cd - > /dev/null
        return
    fi
    
    # 运行迁移
    if command -v alembic &> /dev/null; then
        # 设置数据库URL环境变量
        export DATABASE_URL
        
        # 初始化迁移（如果第一次运行）
        if [ ! -d "alembic/versions" ]; then
            echo -e "${BLUE}初始化Alembic迁移...${NC}"
            alembic init alembic
            # 需要手动配置alembic.ini中的sqlalchemy.url
        fi
        
        # 创建迁移版本
        echo -e "${BLUE}创建迁移版本...${NC}"
        alembic revision --autogenerate -m "初始迁移"
        
        # 应用迁移
        echo -e "${BLUE}应用迁移...${NC}"
        alembic upgrade head
        
        echo -e "${GREEN}✓ 数据库迁移完成${NC}"
    else
        echo -e "${YELLOW}警告: alembic 命令未找到，跳过迁移${NC}"
        echo "请安装: pip install alembic"
    fi
    
    # 返回原目录
    cd - > /dev/null
}

# 函数：创建基础数据
create_base_data() {
    echo -e "${BLUE}创建基础数据...${NC}"
    
    # 检查数据库连接
    if ! check_database_exists; then
        echo -e "${RED}错误: 数据库不存在，无法创建基础数据${NC}"
        return 1
    fi
    
    # 创建基础数据SQL
    BASE_DATA_SQL=$(cat <<EOF
-- 创建系统配置
INSERT INTO system_config (config_key, config_value, description) VALUES
('check_interval', '1800', '自动检查间隔（秒）'),
('manual_cooldown', '300', '手动刷新冷却时间（秒）'),
('history_retention', '90', '历史记录保留天数'),
('ssh_timeout', '30', 'SSH连接超时时间'),
('ssh_max_connections', '10', 'SSH最大连接数'),
('batch_check_size', '100', '批量检查设备数量')
ON CONFLICT (config_key) DO NOTHING;

-- 创建默认角色
INSERT INTO roles (name, description, is_system) VALUES
('super_admin', '超级管理员', true),
('admin', '管理员', true),
('area_admin', '区域管理员', true),
('school_admin', '学校管理员', true),
('operator', '操作员', true),
('viewer', '查看员', true)
ON CONFLICT (name) DO NOTHING;

-- 创建基础权限
INSERT INTO permissions (name, code, module, description) VALUES
-- 设备模块权限
('设备查看', 'device_view', 'device', '查看设备信息'),
('设备编辑', 'device_edit', 'device', '编辑设备信息'),
('设备删除', 'device_delete', 'device', '删除设备'),
('设备导入', 'device_import', 'device', '导入设备数据'),
('设备检查', 'device_check', 'device', '检查设备状态'),
-- 用户模块权限
('用户查看', 'user_view', 'user', '查看用户信息'),
('用户管理', 'user_manage', 'user', '管理用户'),
('角色管理', 'role_manage', 'user', '管理角色'),
('权限管理', 'permission_manage', 'user', '管理权限'),
-- 系统模块权限
('系统配置', 'system_config', 'system', '系统配置管理'),
('OLT管理', 'olt_manage', 'system', 'OLT设备管理'),
('任务管理', 'task_manage', 'system', '任务配置管理'),
('审核管理', 'audit_manage', 'system', '审核流程管理'),
-- 统计模块权限
('统计查看', 'stats_view', 'stats', '查看统计信息'),
('报表导出', 'report_export', 'stats', '导出报表')
ON CONFLICT (code) DO NOTHING;

-- 为超级管理员角色分配所有权限
INSERT INTO role_permissions (role_id, permission_id)
SELECT r.id, p.id
FROM roles r
CROSS JOIN permissions p
WHERE r.name = 'super_admin'
ON CONFLICT DO NOTHING;

-- 为管理员角色分配大部分权限（除了用户管理相关）
INSERT INTO role_permissions (role_id, permission_id)
SELECT r.id, p.id
FROM roles r
CROSS JOIN permissions p
WHERE r.name = 'admin'
  AND p.code NOT IN ('user_manage', 'role_manage', 'permission_manage')
ON CONFLICT DO NOTHING;

-- 为区域管理员分配基础权限
INSERT INTO role_permissions (role_id, permission_id)
SELECT r.id, p.id
FROM roles r
CROSS JOIN permissions p
WHERE r.name = 'area_admin'
  AND p.code IN ('device_view', 'device_edit', 'device_check', 'stats_view')
ON CONFLICT DO NOTHING;

-- 为学校管理员分配基础权限
INSERT INTO role_permissions (role_id, permission_id)
SELECT r.id, p.id
FROM roles r
CROSS JOIN permissions p
WHERE r.name = 'school_admin'
  AND p.code IN ('device_view', 'device_edit', 'device_check')
ON CONFLICT DO NOTHING;

-- 为操作员分配基础权限
INSERT INTO role_permissions (role_id, permission_id)
SELECT r.id, p.id
FROM roles r
CROSS JOIN permissions p
WHERE r.name = 'operator'
  AND p.code IN ('device_view', 'device_edit', 'device_check')
ON CONFLICT DO NOTHING;

-- 为查看员分配只读权限
INSERT INTO role_permissions (role_id, permission_id)
SELECT r.id, p.id
FROM roles r
CROSS JOIN permissions p
WHERE r.name = 'viewer'
  AND p.code IN ('device_view', 'stats_view')
ON CONFLICT DO NOTHING;

-- 创建默认系统用户（通过Casdoor同步，这里只是占位）
INSERT INTO users (casdoor_id, username, email, role, is_active) VALUES
('system', '系统用户', 'system@example.com', 'super_admin', true)
ON CONFLICT (casdoor_id) DO NOTHING;

EOF
)
    
    # 执行SQL
    echo "$BASE_DATA_SQL" | PGPASSWORD="$DB_PASSWORD" psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAME"
    
    if [ $? -eq 0 ]; then
        echo -e "${GREEN}✓ 基础数据创建成功${NC}"
    else
        echo -e "${YELLOW}警告: 基础数据创建可能失败，表可能不存在${NC}"
        echo "请先运行数据库迁移"
    fi
}

# 函数：显示数据库信息
show_database_info() {
    echo -e "\n${GREEN}========================================${NC}"
    echo -e "${GREEN}  数据库初始化完成${NC}"
    echo -e "${GREEN}========================================${NC}"
    echo ""
    echo -e "${BLUE}数据库连接信息:${NC}"
    echo "  主机: $DB_HOST"
    echo "  端口: $DB_PORT"
    echo "  数据库: $DB_NAME"
    echo "  用户: $DB_USER"
    echo ""
    echo -e "${BLUE}连接字符串:${NC}"
    echo "  $DATABASE_URL"
    echo ""
    echo -e "${BLUE}下一步:${NC}"
    echo "  1. 启动后端服务运行数据库迁移"
    echo "  2. 配置Casdoor认证"
    echo "  3. 导入设备数据"
    echo ""
}

# 主程序
main() {
    echo -e "${GREEN}========================================${NC}"
    echo -e "${GREEN}  H3C ONU设备管理系统 - 数据库初始化${NC}"
    echo -e "${GREEN}========================================${NC}"
    echo ""
    
    # 检查PostgreSQL连接
    if ! check_postgres_connection; then
        echo -e "${RED}请检查PostgreSQL服务是否运行，以及连接信息是否正确${NC}"
        exit 1
    fi
    
    # 检查数据库是否存在
    if check_database_exists; then
        echo -e "${YELLOW}数据库已存在，是否重新初始化?${NC}"
        read -p "重新初始化将删除现有数据! (y/n): " -n 1 -r
        echo
        if [[ $REPLY =~ ^[Yy]$ ]]; then
            echo -e "${YELLOW}删除数据库 '$DB_NAME'...${NC}"
            PGPASSWORD="$DB_PASSWORD" psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d postgres -c "DROP DATABASE IF EXISTS $DB_NAME;"
            create_database
        else
            echo -e "${YELLOW}跳过数据库创建${NC}"
        fi
    else
        create_database
    fi
    
    # 创建用户和授予权限
    create_database_user
    grant_privileges
    
    # 运行迁移（如果后端代码已存在）
    run_migrations
    
    # 创建基础数据
    create_base_data
    
    # 显示完成信息
    show_database_info
}

# 执行主程序
main "$@"