#!/usr/bin/env python3
"""Simple MCP server for Portainer that actually works."""

import os
import sys
import json
import asyncio
from typing import Any, Dict, List

# Suppress all logging to stderr
os.environ["MCP_MODE"] = "true"

import mcp.server.stdio
import mcp.types as types
from mcp.server import NotificationOptions, Server
from mcp.server.models import InitializationOptions

# Add src to path
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src'))

# Create server
server = Server("portainer-core")

# Store for our state
portainer_url = os.getenv("PORTAINER_URL", "https://partner.portainer.live")
api_key = os.getenv("PORTAINER_API_KEY", "")


@server.list_tools()
async def handle_list_tools() -> list[types.Tool]:
    """List available tools."""
    return [
        types.Tool(
            name="test_connection",
            description="Test connection to Portainer",
            inputSchema={
                "type": "object",
                "properties": {},
                "required": []
            }
        ),
        types.Tool(
            name="get_users",
            description="Get list of Portainer users",
            inputSchema={
                "type": "object",
                "properties": {},
                "required": []
            }
        ),
        types.Tool(
            name="create_user",
            description="Create a new Portainer user",
            inputSchema={
                "type": "object",
                "properties": {
                    "username": {
                        "type": "string",
                        "description": "Username for the new user"
                    },
                    "password": {
                        "type": "string",
                        "description": "Password for the new user"
                    },
                    "role": {
                        "type": "string",
                        "description": "User role (Administrator, StandardUser, ReadOnlyUser)",
                        "enum": ["Administrator", "StandardUser", "ReadOnlyUser"],
                        "default": "StandardUser"
                    }
                },
                "required": ["username", "password"]
            }
        )
    ]


@server.call_tool()
async def handle_call_tool(name: str, arguments: dict[str, Any] | None) -> list[types.TextContent]:
    """Handle tool calls."""
    import httpx
    
    if name == "test_connection":
        try:
            async with httpx.AsyncClient(verify=False) as client:
                response = await client.get(
                    f"{portainer_url}/api/status",
                    headers={"X-API-Key": api_key} if api_key else {}
                )
                if response.status_code == 200:
                    return [types.TextContent(
                        type="text",
                        text=f"✓ Connected to Portainer at {portainer_url}"
                    )]
                else:
                    return [types.TextContent(
                        type="text",
                        text=f"✗ Failed to connect: HTTP {response.status_code}"
                    )]
        except Exception as e:
            return [types.TextContent(
                type="text",
                text=f"✗ Connection error: {str(e)}"
            )]
    
    elif name == "get_users":
        try:
            async with httpx.AsyncClient(verify=False) as client:
                response = await client.get(
                    f"{portainer_url}/api/users",
                    headers={"X-API-Key": api_key}
                )
                if response.status_code == 200:
                    users = response.json()
                    result = f"Found {len(users)} users:\n"
                    for user in users:
                        # Handle both old (int) and new (string) role formats
                        role = user.get("Role", "Unknown")
                        if isinstance(role, int):
                            role_map = {1: "Administrator", 2: "StandardUser", 3: "ReadOnlyUser"}
                            role = role_map.get(role, f"Unknown({role})")
                        result += f"- {user.get('Username', 'Unknown')} ({role})\n"
                    return [types.TextContent(type="text", text=result)]
                else:
                    return [types.TextContent(
                        type="text",
                        text=f"Failed to get users: HTTP {response.status_code}"
                    )]
        except Exception as e:
            return [types.TextContent(
                type="text",
                text=f"Error getting users: {str(e)}"
            )]
    
    elif name == "create_user":
        if not arguments:
            return [types.TextContent(
                type="text",
                text="Error: Missing required arguments"
            )]
        
        username = arguments.get("username")
        password = arguments.get("password")
        role = arguments.get("role", "StandardUser")
        
        if not username or not password:
            return [types.TextContent(
                type="text",
                text="Error: Username and password are required"
            )]
        
        # Convert role string to integer for older Portainer versions
        role_map = {"Administrator": 1, "StandardUser": 2, "ReadOnlyUser": 3}
        role_int = role_map.get(role, 2)
        
        try:
            async with httpx.AsyncClient(verify=False) as client:
                # Try with integer role first (older versions)
                response = await client.post(
                    f"{portainer_url}/api/users",
                    headers={"X-API-Key": api_key},
                    json={
                        "Username": username,
                        "Password": password,
                        "Role": role_int
                    }
                )
                
                if response.status_code == 409:
                    return [types.TextContent(
                        type="text",
                        text=f"User '{username}' already exists"
                    )]
                elif response.status_code in [200, 201]:
                    return [types.TextContent(
                        type="text",
                        text=f"✓ User '{username}' created successfully with role {role}"
                    )]
                else:
                    # Try with string role (newer versions)
                    response = await client.post(
                        f"{portainer_url}/api/users",
                        headers={"X-API-Key": api_key},
                        json={
                            "Username": username,
                            "Password": password,
                            "Role": role
                        }
                    )
                    if response.status_code in [200, 201]:
                        return [types.TextContent(
                            type="text",
                            text=f"✓ User '{username}' created successfully with role {role}"
                        )]
                    else:
                        return [types.TextContent(
                            type="text",
                            text=f"Failed to create user: HTTP {response.status_code} - {response.text}"
                        )]
        except Exception as e:
            return [types.TextContent(
                type="text",
                text=f"Error creating user: {str(e)}"
            )]
    
    else:
        return [types.TextContent(
            type="text",
            text=f"Unknown tool: {name}"
        )]


async def run():
    """Run the MCP server."""
    # Use stdio transport
    async with mcp.server.stdio.stdio_server() as (read_stream, write_stream):
        await server.run(
            read_stream,
            write_stream,
            InitializationOptions(
                server_name="portainer-core",
                server_version="0.2.0",
                capabilities=server.get_capabilities(
                    notification_options=NotificationOptions(
                        prompts_changed=False,
                        resources_changed=False,
                        tools_changed=False,
                    ),
                    experimental_capabilities={},
                ),
            ),
        )


def main():
    """Main entry point."""
    asyncio.run(run())


if __name__ == "__main__":
    main()