"""
UASS MCP (Model Context Protocol) Server.

This module implements an MCP server that exposes UASS state as tools and resources,
allowing any MCP-compliant AI client to query and interact with agent state.
"""

from typing import Optional, Any
import json
from mcp.server import Server
from mcp.types import Tool, TextContent, Resource, EmbeddedResource

from .client import UASSClient
from .backends import StateBackend


class UASSMCPServer:
    """
    MCP Server for UASS.
    
    Exposes agent state as MCP tools and resources.
    """
    
    def __init__(self, client: UASSClient):
        """
        Initialize the MCP server.
        
        Args:
            client: UASSClient instance connected to a backend
        """
        self.client = client
        self.server = Server("uass-mcp-server")
        self._register_tools()
        self._register_resources()
    
    def _register_tools(self):
        """Register MCP tools."""
        
        @self.server.list_tools()
        async def list_tools() -> list[Tool]:
            """List available tools."""
            return [
                Tool(
                    name="read_agent_state",
                    description="Read the current state of a specific agent in a session",
                    inputSchema={
                        "type": "object",
                        "properties": {
                            "session_id": {
                                "type": "string",
                                "description": "The session identifier"
                            },
                            "agent_id": {
                                "type": "string",
                                "description": "The agent identifier"
                            }
                        },
                        "required": ["session_id", "agent_id"]
                    }
                ),
                Tool(
                    name="list_active_agents",
                    description="List all active agents in a session",
                    inputSchema={
                        "type": "object",
                        "properties": {
                            "session_id": {
                                "type": "string",
                                "description": "The session identifier"
                            }
                        },
                        "required": ["session_id"]
                    }
                ),
                Tool(
                    name="write_agent_state",
                    description="Write or update agent state",
                    inputSchema={
                        "type": "object",
                        "properties": {
                            "session_id": {
                                "type": "string",
                                "description": "The session identifier"
                            },
                            "agent_id": {
                                "type": "string",
                                "description": "The agent identifier"
                            },
                            "framework": {
                                "type": "string",
                                "description": "Framework name (e.g., 'langgraph', 'autogen')"
                            },
                            "variables": {
                                "type": "object",
                                "description": "Agent variables to update"
                            },
                            "context": {
                                "type": "object",
                                "description": "Agent context to update"
                            }
                        },
                        "required": ["session_id", "agent_id", "framework"]
                    }
                ),
                Tool(
                    name="get_shared_state",
                    description="Get shared state for a session (accessible by all agents)",
                    inputSchema={
                        "type": "object",
                        "properties": {
                            "session_id": {
                                "type": "string",
                                "description": "The session identifier"
                            }
                        },
                        "required": ["session_id"]
                    }
                )
            ]
        
        @self.server.call_tool()
        async def call_tool(name: str, arguments: dict) -> list[TextContent]:
            """Handle tool calls."""
            
            if name == "read_agent_state":
                session_id = arguments["session_id"]
                agent_id = arguments["agent_id"]
                
                state = self.client.read_state(session_id, agent_id)
                
                if state is None:
                    return [TextContent(
                        type="text",
                        text=f"No state found for agent '{agent_id}' in session '{session_id}'"
                    )]
                
                return [TextContent(
                    type="text",
                    text=json.dumps(state, indent=2)
                )]
            
            elif name == "list_active_agents":
                session_id = arguments["session_id"]
                agents = self.client.list_agents(session_id)
                
                return [TextContent(
                    type="text",
                    text=json.dumps({
                        "session_id": session_id,
                        "agents": agents,
                        "count": len(agents)
                    }, indent=2)
                )]
            
            elif name == "write_agent_state":
                session_id = arguments["session_id"]
                agent_id = arguments["agent_id"]
                framework = arguments["framework"]
                variables = arguments.get("variables", {})
                context = arguments.get("context", {})
                
                success = self.client.write_state(
                    session_id=session_id,
                    agent_id=agent_id,
                    framework=framework,
                    variables=variables,
                    context=context
                )
                
                return [TextContent(
                    type="text",
                    text=json.dumps({
                        "success": success,
                        "message": f"State {'updated' if success else 'update failed'} for agent '{agent_id}'"
                    }, indent=2)
                )]
            
            elif name == "get_shared_state":
                session_id = arguments["session_id"]
                shared_state = self.client.read_shared_state(session_id)
                
                if shared_state is None:
                    return [TextContent(
                        type="text",
                        text=f"No shared state found for session '{session_id}'"
                    )]
                
                return [TextContent(
                    type="text",
                    text=json.dumps(shared_state, indent=2)
                )]
            
            else:
                return [TextContent(
                    type="text",
                    text=f"Unknown tool: {name}"
                )]
    
    def _register_resources(self):
        """Register MCP resources."""
        
        @self.server.list_resources()
        async def list_resources() -> list[Resource]:
            """List available resources."""
            return [
                Resource(
                    uri="uass://sessions",
                    name="UASS Sessions",
                    description="List of all active UASS sessions",
                    mimeType="application/json"
                )
            ]
        
        @self.server.read_resource()
        async def read_resource(uri: str) -> str:
            """Read a resource by URI."""
            
            # Convert AnyUrl to string if needed
            uri_str = str(uri)
            
            if uri_str == "uass://sessions":
                # This is a placeholder - in a real implementation,
                # we'd need to track sessions in the backend
                return json.dumps({
                    "message": "Session listing not yet implemented",
                    "hint": "Use list_active_agents tool with a known session_id"
                }, indent=2)
            
            # Parse URI like: uass://sessions/{session_id}/agents/{agent_id}
            if uri_str.startswith("uass://sessions/"):
                parts = uri_str.replace("uass://sessions/", "").split("/")
                
                if len(parts) >= 3 and parts[1] == "agents":
                    session_id = parts[0]
                    agent_id = parts[2]
                    
                    state = self.client.read_state(session_id, agent_id)
                    
                    if state is None:
                        return json.dumps({
                            "error": f"No state found for agent '{agent_id}' in session '{session_id}'"
                        }, indent=2)
                    
                    return json.dumps(state, indent=2)
            
            return json.dumps({
                "error": f"Unknown resource URI: {uri_str}"
            }, indent=2)
    
    def get_server(self) -> Server:
        """Get the underlying MCP server instance."""
        return self.server
