from fastmcp import FastMCP
import sqlite3
from typing import List, Dict

# 创建MCP服务器
mcp = FastMCP("企业用户管理系统")

# 数据库初始化（生产环境使用PostgreSQL）
def init_database():
    conn = sqlite3.connect('users.db')
    cursor = conn.cursor()
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS users (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            username TEXT UNIQUE NOT NULL,
            email TEXT UNIQUE NOT NULL,
            name TEXT NOT NULL,
            status TEXT DEFAULT 'active',
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
    ''')
    conn.commit()
    conn.close()

init_database()  # 启动时初始化

@mcp.tool
def get_user_by_id(user_id: int) -> Dict:
    """根据ID查询用户信息"""
    conn = sqlite3.connect('users.db')
    cursor = conn.cursor()
    cursor.execute(
        'SELECT id, username, email, name, status FROM users WHERE id = ?',
        (user_id,)
    )
    result = cursor.fetchone()
    conn.close()
    
    if result:
        return {
            'id': result[0], 'username': result[1], 'email': result[2],
            'name': result[3], 'status': result[4]
        }
    else:
        raise ValueError(f"未找到ID为{user_id}的用户")

@mcp.tool
def create_user(username: str, email: str, name: str) -> Dict:
    """创建新用户"""
    conn = sqlite3.connect('users.db')
    cursor = conn.cursor()
    
    try:
        cursor.execute(
            'INSERT INTO users (username, email, name) VALUES (?, ?, ?)',
            (username, email, name)
        )
        user_id = cursor.lastrowid
        conn.commit()
        return {"id": user_id, "username": username, "status": "created"}
    except sqlite3.IntegrityError:
        raise ValueError("用户名或邮箱已存在")
    finally:
        conn.close()

@mcp.resource("users://stats")
def get_user_stats() -> Dict:
    """获取用户统计信息"""
    conn = sqlite3.connect('users.db')
    cursor = conn.cursor()
    
    # 统计活跃用户数
    cursor.execute('SELECT COUNT(*) FROM users WHERE status = "active"')
    active_count = cursor.fetchone()[0]
    
    # 统计总用户数
    cursor.execute('SELECT COUNT(*) FROM users')
    total_count = cursor.fetchone()[0]
    
    conn.close()
    
    return {
        "active_users": active_count,
        "total_users": total_count,
        "last_updated": "2025-01-15 10:30:00"
    }

if __name__ == "__main__":
    print("🚀 企业用户管理系统已启动")
    mcp.run()

@mcp.tool
def send_welcome_email(user_id: int) -> Dict:
    """给新用户发送欢迎邮件"""
    user = get_user_by_id(user_id)
    
    # 邮件数据
    email_data = {
        "to": user['email'],
        "subject": f"欢迎加入，{user['name']}！",
        "template": "welcome",
        "variables": {"name": user['name']}
    }
    
    # 调用邮件服务API（生产环境替换为真实API）
    try:
        # response = requests.post("https://api.emailservice.com/send", json=email_data)
        return {
            "status": "sent",
            "recipient": user['email'],
            "sent_at": "2025-01-15 10:30:00"
        }
    except Exception as e:
        raise ValueError(f"邮件发送失败：{str(e)}")

@mcp.tool
def validate_email_domain(email: str) -> Dict:
    """验证邮箱域名安全性"""
    domain = email.split('@')[-1]
    
    # 模拟域名验证API调用
    trusted_domains = ['gmail.com', 'outlook.com', '163.com', 'qq.com']
    is_trusted = domain in trusted_domains
    
    return {
        "email": email,
        "domain": domain,
        "is_trusted": is_trusted,
        "risk_level": "low" if is_trusted else "medium"
    }

@mcp.tool
def export_users_csv() -> Dict:
    """导出用户数据到CSV文件"""
    import csv
    from datetime import datetime
    
    filename = f"users_export_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv"
    
    # 获取所有用户数据
    conn = sqlite3.connect('users.db')
    cursor = conn.cursor()
    cursor.execute('SELECT id, username, email, name, status FROM users')
    users = cursor.fetchall()
    conn.close()
    
    # 写入CSV
    with open(filename, 'w', newline='', encoding='utf-8') as csvfile:
        writer = csv.writer(csvfile)
        writer.writerow(['ID', '用户名', '邮箱', '姓名', '状态'])
        writer.writerows(users)
    
    return {
        "filename": filename,
        "record_count": len(users),
        "exported_at": datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    }

# test_system.py - 快速验证
import asyncio
from fastmcp import Client
from user_manager import mcp

async def test_core_features():
    async with Client(mcp) as client:
        # 创建测试用户
        user = await client.call_tool("create_user", {
            "username": "testuser", "email": "test@example.com", "name": "测试用户"
        })
        print(f"✅ 用户创建成功: {user.data}")
        
        # 查询用户信息
        user_info = await client.call_tool("get_user_by_id", {"user_id": 1})
        print(f"✅ 用户查询成功: {user_info.data}")
        
        # 获取统计数据
        stats = await client.read_resource("users://stats")
        print(f"✅ 统计数据获取: {stats.contents[0].text}")

if __name__ == "__main__":
    asyncio.run(test_core_features())

# 1. 输入验证 - 数据安全第一道防线
def validate_email(email: str) -> bool:
    import re
    pattern = r'^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}$'
    return re.match(pattern, email) is not None

@mcp.tool
def safe_create_user(username: str, email: str, name: str) -> Dict:
    """安全的用户创建（带验证）"""
    if not validate_email(email):
        raise ValueError("邮箱格式不正确")
    
    if len(username) < 3:
        raise ValueError("用户名至少3个字符")
    
    return create_user(username, email, name)

# 2. 环境变量配置 - 敏感信息不写死
import os
DATABASE_URL = os.getenv("DATABASE_URL", "sqlite:///users.db")
API_KEY = os.getenv("EMAIL_API_KEY", "your-api-key")

# 连接池管理 - 提升数据库性能
from functools import lru_cache

@lru_cache(maxsize=128)
def cached_user_query(user_id: int, timestamp: int):
    """带缓存的用户查询（5分钟缓存）"""
    return get_user_by_id(user_id)

@mcp.tool
def fast_user_query(user_id: int) -> Dict:
    """快速用户查询"""
    import time
    cache_key = int(time.time() // 300)  # 5分钟缓存
    return cached_user_query(user_id, cache_key)