"""
Custom Schema and Tool Definition Example

Demonstrates how to define custom MCP tools with detailed schemas,
input validation, and output formatting.
"""

from typing import Optional, List
from pydantic import BaseModel, Field
from drf_mcp.schema import SchemaConverter
from drf_mcp.provider import MCPProvider


# Define input/output models using Pydantic
class UserFilterInput(BaseModel):
    """Input parameters for user filtering."""
    username: Optional[str] = Field(
        None,
        description="Filter users by username (supports partial match)"
    )
    is_active: Optional[bool] = Field(
        None,
        description="Filter by active status"
    )
    email_domain: Optional[str] = Field(
        None,
        description="Filter users by email domain (e.g., '@example.com')"
    )


class UserOutput(BaseModel):
    """Structured user output."""
    id: int = Field(description="User ID")
    username: str = Field(description="Username")
    email: str = Field(description="Email address")
    is_active: bool = Field(description="Whether user is active")
    is_staff: bool = Field(description="Whether user is staff member")


class UsersListOutput(BaseModel):
    """List of users with metadata."""
    count: int = Field(description="Total number of users")
    results: List[UserOutput] = Field(description="List of user objects")


# Create custom MCP tool
async def search_users(params: UserFilterInput) -> UsersListOutput:
    """
    Search for users based on filtering criteria.
    
    This tool demonstrates:
    - Structured input validation with Pydantic
    - Async execution
    - Detailed output formatting
    - Clear documentation for LLM context
    
    Args:
        params: User search parameters
        
    Returns:
        List of matching users with pagination info
        
    Example:
        To find active users with 'john' in username:
        - username: "john"
        - is_active: true
    """
    from django.contrib.auth.models import User
    
    # Build query
    queryset = User.objects.all()
    
    if params.username:
        queryset = queryset.filter(username__icontains=params.username)
    
    if params.is_active is not None:
        queryset = queryset.filter(is_active=params.is_active)
    
    if params.email_domain:
        queryset = queryset.filter(email__endswith=params.email_domain)
    
    # Format results
    users = [
        UserOutput(
            id=user.id,
            username=user.username,
            email=user.email,
            is_active=user.is_active,
            is_staff=user.is_staff
        )
        for user in queryset
    ]
    
    return UsersListOutput(
        count=len(users),
        results=users
    )


# Register custom tool with MCP Provider
provider = MCPProvider()

# Define tool schema manually
custom_tool_schema = {
    'name': 'search_users',
    'description': 'Search for users based on filtering criteria',
    'inputSchema': {
        'type': 'object',
        'properties': {
            'username': {
                'type': 'string',
                'description': 'Filter users by username (supports partial match)'
            },
            'is_active': {
                'type': 'boolean',
                'description': 'Filter by active status'
            },
            'email_domain': {
                'type': 'string',
                'description': 'Filter users by email domain'
            }
        }
    }
}

provider.register_custom_tool(
    name='search_users',
    schema=custom_tool_schema,
    handler=search_users
)


# Usage in MCP Server
if __name__ == "__main__":
    from fastmcp import Server
    import asyncio
    
    mcp_server = Server("django-mcp-custom")
    
    # Register custom tools
    mcp_server.tool(search_users)
    
    # Test the tool
    async def test():
        result = await search_users(
            UserFilterInput(username="admin", is_active=True)
        )
        print(f"Found {result.count} users")
        for user in result.results:
            print(f"  - {user.username} ({user.email})")
    
    # Uncomment to test: asyncio.run(test())
    print("Custom schema tool registered")
