"""
Real-world example: AI Agent Tool with MCP Guardian Security

This example demonstrates how MCP Guardian can be used in a real AI Agent scenario
where different agents have different permissions for accessing tools.
"""

from fastapi import FastAPI, Request, HTTPException
from mcp_guardian import Guardian
import json

# Initialize FastAPI app
app = FastAPI(
    title="AI Agent Tools API",
    description="Secure API for AI Agents to access various tools",
    version="1.0.0"
)

# Initialize Guardian with configuration
# Note: Copy config.example.yaml to config.yaml and update with your settings
guardian = Guardian(config_path="config.yaml")

# Public health check endpoint
@app.get("/health")
async def health_check():
    """Public health check endpoint."""
    return {
        "status": "healthy",
        "service": "AI Agent Tools API",
        "version": "1.0.0",
        "security": "MCP Guardian enabled"
    }

# Calculator tool - requires calculator scope
@app.post("/tools/calculator")
@guardian.protect(scope="tools:calculator:execute")
async def calculator(request: Request, expression: str = "2 + 2"):
    """
    Calculator tool for AI agents.
    
    Requires: tools:calculator:execute scope
    Rate limit: Per user (from config)
    """
    user_info = request.state.token
    
    try:
        # Simple and safe expression evaluation
        allowed_chars = set('0123456789+-*/(). ')
        if not all(c in allowed_chars for c in expression):
            raise ValueError("Invalid characters in expression")
        
        result = eval(expression)
        
        return {
            "tool": "calculator",
            "expression": expression,
            "result": result,
            "user": user_info.get("sub"),
            "timestamp": user_info.get("iat")
        }
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Calculation error: {str(e)}")

# Weather tool - requires weather scope
@app.get("/tools/weather")
@guardian.protect(scope="tools:weather:read")
async def weather(request: Request, city: str = "New York"):
    """
    Weather information tool for AI agents.
    
    Requires: tools:weather:read scope
    Rate limit: Per user (from config)
    """
    user_info = request.state.token
    
    # Mock weather data (in real app, this would call a weather API)
    weather_data = {
        "city": city,
        "temperature": "72°F",
        "condition": "Sunny",
        "humidity": "65%",
        "wind": "5 mph"
    }
    
    return {
        "tool": "weather",
        "data": weather_data,
        "user": user_info.get("sub"),
        "timestamp": user_info.get("iat")
    }

# File system tool - requires file access scope
@app.get("/tools/files")
@guardian.protect(scope="tools:files:read")
async def list_files(request: Request, path: str = "/"):
    """
    File system listing tool for AI agents.
    
    Requires: tools:files:read scope
    Rate limit: Per user (from config)
    """
    user_info = request.state.token
    
    # Mock file listing (in real app, this would list actual files)
    files = [
        {"name": "document1.txt", "size": "1.2KB", "modified": "2024-01-15"},
        {"name": "image1.jpg", "size": "2.5MB", "modified": "2024-01-14"},
        {"name": "data.json", "size": "856B", "modified": "2024-01-13"}
    ]
    
    return {
        "tool": "files",
        "path": path,
        "files": files,
        "user": user_info.get("sub"),
        "timestamp": user_info.get("iat")
    }

# Admin tool - requires admin scope
@app.post("/tools/admin/system")
@guardian.protect(scope="tools:admin:system")
async def system_info(request: Request):
    """
    System information tool for AI agents with admin access.
    
    Requires: tools:admin:system scope
    Rate limit: Per user (from config)
    """
    user_info = request.state.token
    
    # Mock system info (in real app, this would return actual system data)
    system_info = {
        "cpu_usage": "45%",
        "memory_usage": "2.1GB / 8GB",
        "disk_usage": "45GB / 100GB",
        "uptime": "7 days, 3 hours",
        "active_users": 15
    }
    
    return {
        "tool": "admin_system",
        "data": system_info,
        "user": user_info.get("sub"),
        "timestamp": user_info.get("iat"),
        "admin_access": True
    }

# User info endpoint - shows current user's permissions
@app.get("/user/info")
@guardian.protect(scope="user:info:read")
async def user_info(request: Request):
    """
    Get current user information and permissions.
    
    Requires: user:info:read scope
    """
    user_info = request.state.token
    
    return {
        "user_id": user_info.get("sub"),
        "scopes": user_info.get("scopes", []),
        "audience": user_info.get("aud"),
        "issued_at": user_info.get("iat"),
        "expires_at": user_info.get("exp"),
        "issuer": user_info.get("iss")
    }

# Rate limit info endpoint
@app.get("/user/rate-limit")
@guardian.protect(scope="user:rate-limit:read")
async def rate_limit_info(request: Request):
    """
    Get current user's rate limit information.
    
    Requires: user:rate-limit:read scope
    """
    user_info = request.state.token
    
    # In a real implementation, you'd query the rate limiter
    return {
        "user_id": user_info.get("sub"),
        "rate_limit": {
            "requests": 100,
            "period_minutes": 60,
            "current_usage": "Not implemented in demo"
        }
    }

if __name__ == "__main__":
    import uvicorn
    print("🚀 Starting AI Agent Tools API with MCP Guardian Security")
    print("📋 Available tools:")
    print("  - Calculator: tools:calculator:execute")
    print("  - Weather: tools:weather:read")
    print("  - Files: tools:files:read")
    print("  - Admin: tools:admin:system")
    print("  - User Info: user:info:read")
    print("  - Rate Limit: user:rate-limit:read")
    print()
    print("🔐 Security: All endpoints protected with JWT authentication and scope-based authorization")
    print("⚡ Rate Limiting: Per-user rate limiting enabled")
    print()
    
    uvicorn.run(app, host="0.0.0.0", port=8000)
