#!/bin/bash

# TJWater Server 安全功能快速设置脚本

set -e

echo "=================================="
echo "TJWater Server 安全功能设置"
echo "=================================="
echo ""

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

# 步骤 1: 检查依赖
echo "📦 步骤 1/5: 检查 Python 依赖..."
if ! python -c "import cryptography, passlib, jose" 2>/dev/null; then
    echo -e "${YELLOW}缺少依赖，正在安装...${NC}"
    pip install cryptography passlib python-jose bcrypt
else
    echo -e "${GREEN}✓ 依赖已安装${NC}"
fi
echo ""

# 步骤 2: 生成密钥
echo "🔑 步骤 2/5: 生成安全密钥..."

if [ ! -f .env ]; then
    echo "正在创建 .env 文件..."
    cp .env.example .env
    
    # 生成 JWT 密钥
    JWT_KEY=$(openssl rand -hex 32)
    sed -i "s/SECRET_KEY=.*/SECRET_KEY=$JWT_KEY/" .env
    echo -e "${GREEN}✓ JWT 密钥已生成${NC}"
    
    # 生成加密密钥
    ENC_KEY=$(python -c "from cryptography.fernet import Fernet; print(Fernet.generate_key().decode())")
    sed -i "s/ENCRYPTION_KEY=.*/ENCRYPTION_KEY=$ENC_KEY/" .env
    echo -e "${GREEN}✓ 加密密钥已生成${NC}"
else
    echo -e "${YELLOW}⚠ .env 文件已存在，跳过生成${NC}"
fi
echo ""

# 步骤 3: 数据库配置
echo "💾 步骤 3/5: 数据库配置..."
read -p "请输入数据库名称 [默认: tjwater]: " DB_NAME
DB_NAME=${DB_NAME:-tjwater}

read -p "请输入数据库用户 [默认: postgres]: " DB_USER
DB_USER=${DB_USER:-postgres}

read -sp "请输入数据库密码: " DB_PASS
echo ""

# 更新 .env
sed -i "s/DB_NAME=.*/DB_NAME=$DB_NAME/" .env
sed -i "s/DB_USER=.*/DB_USER=$DB_USER/" .env
sed -i "s/DB_PASSWORD=.*/DB_PASSWORD=$DB_PASS/" .env

echo -e "${GREEN}✓ 数据库配置已更新${NC}"
echo ""

# 步骤 4: 执行数据库迁移
echo "🗄️  步骤 4/5: 执行数据库迁移..."
read -p "是否立即执行数据库迁移？(y/n) [y]: " DO_MIGRATION
DO_MIGRATION=${DO_MIGRATION:-y}

if [ "$DO_MIGRATION" = "y" ]; then
    echo "正在执行迁移脚本..."
    
    PGPASSWORD=$DB_PASS psql -U $DB_USER -d $DB_NAME -f resources/sql/001_create_users_table.sql 2>&1 | grep -v "NOTICE"
    if [ $? -eq 0 ]; then
        echo -e "${GREEN}✓ 用户表创建成功${NC}"
    else
        echo -e "${RED}✗ 用户表创建失败${NC}"
    fi
    
    PGPASSWORD=$DB_PASS psql -U $DB_USER -d $DB_NAME -f resources/sql/002_create_audit_logs_table.sql 2>&1 | grep -v "NOTICE"
    if [ $? -eq 0 ]; then
        echo -e "${GREEN}✓ 审计日志表创建成功${NC}"
    else
        echo -e "${RED}✗ 审计日志表创建失败${NC}"
    fi
else
    echo -e "${YELLOW}⚠ 跳过数据库迁移，请稍后手动执行：${NC}"
    echo "  psql -U $DB_USER -d $DB_NAME -f resources/sql/001_create_users_table.sql"
    echo "  psql -U $DB_USER -d $DB_NAME -f resources/sql/002_create_audit_logs_table.sql"
fi
echo ""

# 步骤 5: 测试
echo "🧪 步骤 5/5: 运行测试..."
if python tests/test_encryption.py 2>&1; then
    echo -e "${GREEN}✓ 加密功能测试通过${NC}"
else
    echo -e "${RED}✗ 加密功能测试失败${NC}"
fi
echo ""

# 完成
echo "=================================="
echo -e "${GREEN}✅ 设置完成！${NC}"
echo "=================================="
echo ""
echo "默认管理员账号："
echo "  用户名: admin"
echo "  密码: admin123"
echo ""
echo "  用户名: tjwater"
echo "  密码: tjwater@123"
echo ""
echo "下一步："
echo "  1. 查看文档: cat SECURITY_README.md"
echo "  2. 查看部署指南: cat DEPLOYMENT.md"
echo "  3. 启动服务器: uvicorn app.main:app --reload"
echo "  4. 访问文档: http://localhost:8000/docs"
echo ""