"""
MCP Server for Spend Analysis - Core Protocol Implementation
"""
import json
import asyncio
from typing import Dict, List, Any, Optional
from dataclasses import dataclass
from enum import Enum
import logging

# MCP Protocol Types
class MessageType(Enum):
    REQUEST = "request"
    RESPONSE = "response"
    NOTIFICATION = "notification"

@dataclass
class MCPMessage:
    jsonrpc: str = "2.0"
    id: Optional[str] = None
    method: Optional[str] = None
    params: Optional[Dict] = None
    result: Optional[Any] = None
    error: Optional[Dict] = None

class MCPServer:
    def __init__(self):
        self.tools = {}
        self.resources = {}
        self.prompts = {}
        self.logger = logging.getLogger(__name__)
        
    def register_tool(self, name: str, description: str, handler):
        """Register a tool that Claude can call"""
        self.tools[name] = {
            "description": description,
            "handler": handler,
            "input_schema": {
                "type": "object",
                "properties": {},
                "required": []
            }
        }
    
    def register_resource(self, uri: str, name: str, description: str, handler):
        """Register a resource that provides data"""
        self.resources[uri] = {
            "name": name,
            "description": description,
            "handler": handler,
            "mimeType": "application/json"
        }
    
    async def handle_message(self, message: Dict) -> Dict:
        """Handle incoming MCP messages"""
        try:
            method = message.get("method")
            params = message.get("params", {})
            msg_id = message.get("id")
            
            if method == "initialize":
                return self._handle_initialize(msg_id)
            elif method == "tools/list":
                return self._handle_list_tools(msg_id)
            elif method == "tools/call":
                return await self._handle_call_tool(msg_id, params)
            elif method == "resources/list":
                return self._handle_list_resources(msg_id)
            elif method == "resources/read":
                return await self._handle_read_resource(msg_id, params)
            else:
                return self._error_response(msg_id, -32601, f"Method not found: {method}")
                
        except Exception as e:
            self.logger.error(f"Error handling message: {e}")
            return self._error_response(message.get("id"), -32603, str(e))
    
    def _handle_initialize(self, msg_id: str) -> Dict:
        """Handle MCP initialization"""
        return {
            "jsonrpc": "2.0",
            "id": msg_id,
            "result": {
                "protocolVersion": "2024-11-05",
                "capabilities": {
                    "tools": {},
                    "resources": {},
                    "prompts": {}
                },
                "serverInfo": {
                    "name": "spend-analyzer-mcp",
                    "version": "1.0.0"
                }
            }
        }
    
    def _handle_list_tools(self, msg_id: str) -> Dict:
        """List available tools"""
        tools_list = []
        for name, tool in self.tools.items():
            tools_list.append({
                "name": name,
                "description": tool["description"],
                "inputSchema": tool["input_schema"]
            })
        
        return {
            "jsonrpc": "2.0",
            "id": msg_id,
            "result": {"tools": tools_list}
        }
    
    async def _handle_call_tool(self, msg_id: str, params: Dict) -> Dict:
        """Execute a tool call"""
        tool_name = params.get("name")
        arguments = params.get("arguments", {})
        
        if tool_name not in self.tools:
            return self._error_response(msg_id, -32602, f"Tool not found: {tool_name}")
        
        try:
            handler = self.tools[tool_name]["handler"]
            result = await handler(arguments)
            return {
                "jsonrpc": "2.0",
                "id": msg_id,
                "result": {
                    "content": [
                        {
                            "type": "text",
                            "text": json.dumps(result)
                        }
                    ]
                }
            }
        except Exception as e:
            return self._error_response(msg_id, -32603, f"Tool execution failed: {str(e)}")
    
    def _handle_list_resources(self, msg_id: str) -> Dict:
        """List available resources"""
        resources_list = []
        for uri, resource in self.resources.items():
            resources_list.append({
                "uri": uri,
                "name": resource["name"],
                "description": resource["description"],
                "mimeType": resource["mimeType"]
            })
        
        return {
            "jsonrpc": "2.0",
            "id": msg_id,
            "result": {"resources": resources_list}
        }
    
    async def _handle_read_resource(self, msg_id: str, params: Dict) -> Dict:
        """Read a resource"""
        uri = params.get("uri")
        
        if uri not in self.resources:
            return self._error_response(msg_id, -32602, f"Resource not found: {uri}")
        
        try:
            handler = self.resources[uri]["handler"]
            content = await handler()
            return {
                "jsonrpc": "2.0",
                "id": msg_id,
                "result": {
                    "contents": [
                        {
                            "uri": uri,
                            "mimeType": "application/json",
                            "text": json.dumps(content, indent=2)
                        }
                    ]
                }
            }
        except Exception as e:
            return self._error_response(msg_id, -32603, f"Resource read failed: {str(e)}")
    
    def _error_response(self, msg_id: str, code: int, message: str) -> Dict:
        """Create error response"""
        return {
            "jsonrpc": "2.0",
            "id": msg_id,
            "error": {
                "code": code,
                "message": message
            }
        }

# Example usage and testing
if __name__ == "__main__":
    # Test the MCP server
    server = MCPServer()
    
    # Register a simple tool
    async def test_tool(args):
        return f"Test tool called with: {args}"
    
    server.register_tool("test", "A test tool", test_tool)
    
    # Test message handling
    async def test_server():
        init_msg = {
            "jsonrpc": "2.0",
            "id": "1",
            "method": "initialize",
            "params": {}
        }
        
        response = await server.handle_message(init_msg)
        print("Initialize response:", json.dumps(response, indent=2))
        
        list_tools_msg = {
            "jsonrpc": "2.0",
            "id": "2",
            "method": "tools/list"
        }
        
        response = await server.handle_message(list_tools_msg)
        print("List tools response:", json.dumps(response, indent=2))
    
    asyncio.run(test_server())