"""
QuinkGL MCP Server

Provides MCP (Model Context Protocol) interface for monitoring
gossip learning experiments.

Usage:
    # Start server from command line
    python -m quinkgl.mcp.server --metrics-dir ./logs
    
    # Or programmatically
    from quinkgl.mcp import MCPServer, MetricsCollector
    
    collector = MetricsCollector()
    server = MCPServer(collector)
    await server.start()
"""

import asyncio
import json
import sys
import logging
from typing import Optional, Any, Callable, Dict, List
from dataclasses import dataclass

from quinkgl.mcp.metrics import MetricsCollector

logger = logging.getLogger(__name__)


@dataclass
class Tool:
    """MCP Tool definition."""
    name: str
    description: str
    input_schema: dict
    handler: Callable


class MCPServer:
    """
    MCP Server for QuinkGL monitoring.
    
    Exposes tools for AI assistants to query gossip learning metrics.
    
    Available Tools:
        - get_nodes: List all monitored nodes
        - get_training_progress: Get training progress for all nodes
        - get_accuracy_history: Get accuracy over rounds for a node
        - get_model_exchanges: Get model send/receive records
        - get_network_topology: Get peer connection graph
        - analyze_convergence: Analyze model convergence
        - get_full_report: Get complete experiment report
    
    Example:
        collector = MetricsCollector()
        collector.attach(node)
        
        server = MCPServer(collector)
        await server.start()
    """
    
    PROTOCOL_VERSION = "2024-11-05"
    SERVER_NAME = "quinkgl-monitor"
    SERVER_VERSION = "0.1.0"
    
    def __init__(self, metrics_collector: MetricsCollector):
        """
        Initialize MCP Server.
        
        Args:
            metrics_collector: MetricsCollector instance with attached nodes
        """
        self.collector = metrics_collector
        self._tools: Dict[str, Tool] = {}
        self._running = False
        
        # Register default tools
        self._register_default_tools()
    
    def _register_default_tools(self) -> None:
        """Register the default MCP tools."""
        
        # Tool: get_nodes
        self.register_tool(
            name="get_nodes",
            description="List all monitored gossip learning nodes with their current status",
            input_schema={
                "type": "object",
                "properties": {},
                "required": []
            },
            handler=self._handle_get_nodes
        )
        
        # Tool: get_training_progress
        self.register_tool(
            name="get_training_progress",
            description="Get training progress for all nodes including current round, accuracy, and loss",
            input_schema={
                "type": "object",
                "properties": {
                    "node_id": {
                        "type": "string",
                        "description": "Optional: specific node ID to query"
                    }
                },
                "required": []
            },
            handler=self._handle_get_training_progress
        )
        
        # Tool: get_accuracy_history
        self.register_tool(
            name="get_accuracy_history",
            description="Get accuracy and loss history for each training round",
            input_schema={
                "type": "object",
                "properties": {
                    "node_id": {
                        "type": "string",
                        "description": "Optional: specific node ID. If not provided, returns all nodes"
                    }
                },
                "required": []
            },
            handler=self._handle_get_accuracy_history
        )
        
        # Tool: get_model_exchanges
        self.register_tool(
            name="get_model_exchanges",
            description="Get records of model weights sent and received between peers",
            input_schema={
                "type": "object",
                "properties": {
                    "node_id": {
                        "type": "string",
                        "description": "Optional: filter by node ID"
                    },
                    "direction": {
                        "type": "string",
                        "enum": ["sent", "received", "all"],
                        "description": "Filter by direction (default: all)"
                    }
                },
                "required": []
            },
            handler=self._handle_get_model_exchanges
        )
        
        # Tool: get_network_topology
        self.register_tool(
            name="get_network_topology",
            description="Get the network topology showing which nodes communicated with each other",
            input_schema={
                "type": "object",
                "properties": {},
                "required": []
            },
            handler=self._handle_get_network_topology
        )
        
        # Tool: analyze_convergence
        self.register_tool(
            name="analyze_convergence",
            description="Analyze model convergence across the gossip learning network",
            input_schema={
                "type": "object",
                "properties": {},
                "required": []
            },
            handler=self._handle_analyze_convergence
        )
        
        # Tool: get_full_report
        self.register_tool(
            name="get_full_report",
            description="Get a complete report of the gossip learning experiment",
            input_schema={
                "type": "object",
                "properties": {},
                "required": []
            },
            handler=self._handle_get_full_report
        )
    
    def register_tool(
        self,
        name: str,
        description: str,
        input_schema: dict,
        handler: Callable
    ) -> None:
        """Register a new MCP tool."""
        self._tools[name] = Tool(
            name=name,
            description=description,
            input_schema=input_schema,
            handler=handler
        )
    
    # =========================================================================
    # Tool Handlers
    # =========================================================================
    
    async def _handle_get_nodes(self, arguments: dict) -> dict:
        """Handle get_nodes tool call."""
        nodes = self.collector.get_all_metrics()
        return {
            "total_nodes": len(nodes),
            "nodes": [
                {
                    "node_id": node_id,
                    "domain": data["domain"],
                    "is_running": data["is_running"],
                    "current_round": data["current_round"],
                    "final_accuracy": data.get("final_accuracy")
                }
                for node_id, data in nodes.items()
            ]
        }
    
    async def _handle_get_training_progress(self, arguments: dict) -> dict:
        """Handle get_training_progress tool call."""
        node_id = arguments.get("node_id")
        
        if node_id:
            metrics = self.collector.get_node_metrics(node_id)
            if not metrics:
                return {"error": f"Node '{node_id}' not found"}
            
            latest = metrics.training_history[-1] if metrics.training_history else None
            return {
                "node_id": node_id,
                "current_round": metrics.current_round,
                "is_running": metrics.is_running,
                "latest_accuracy": latest.accuracy if latest else None,
                "latest_loss": latest.loss if latest else None,
                "total_training_rounds": len(metrics.training_history),
                "final_accuracy": metrics.final_accuracy
            }
        
        # Return all nodes
        result = {}
        for nid in self.collector.get_all_nodes():
            metrics = self.collector.get_node_metrics(nid)
            if metrics:
                latest = metrics.training_history[-1] if metrics.training_history else None
                result[nid] = {
                    "current_round": metrics.current_round,
                    "is_running": metrics.is_running,
                    "latest_accuracy": latest.accuracy if latest else None,
                    "latest_loss": latest.loss if latest else None,
                    "total_training_rounds": len(metrics.training_history)
                }
        
        return {"nodes": result}
    
    async def _handle_get_accuracy_history(self, arguments: dict) -> dict:
        """Handle get_accuracy_history tool call."""
        node_id = arguments.get("node_id")
        
        if node_id:
            history = self.collector.get_accuracy_history(node_id)
            if not history:
                return {"error": f"No data for node '{node_id}'"}
            return {
                "node_id": node_id,
                "history": history
            }
        
        # Return all nodes
        return {
            "nodes": self.collector.get_all_accuracy_histories()
        }
    
    async def _handle_get_model_exchanges(self, arguments: dict) -> dict:
        """Handle get_model_exchanges tool call."""
        node_id = arguments.get("node_id")
        direction = arguments.get("direction", "all")
        
        exchanges = self.collector.get_model_exchanges(node_id)
        
        if direction != "all":
            exchanges = [ex for ex in exchanges if ex["direction"] == direction]
        
        return {
            "total_exchanges": len(exchanges),
            "exchanges": exchanges
        }
    
    async def _handle_get_network_topology(self, arguments: dict) -> dict:
        """Handle get_network_topology tool call."""
        return self.collector.get_network_topology()
    
    async def _handle_analyze_convergence(self, arguments: dict) -> dict:
        """Handle analyze_convergence tool call."""
        return self.collector.get_convergence_analysis()
    
    async def _handle_get_full_report(self, arguments: dict) -> dict:
        """Handle get_full_report tool call."""
        return self.collector.export_json()
    
    # =========================================================================
    # MCP Protocol Implementation
    # =========================================================================
    
    async def _read_message(self) -> Optional[dict]:
        """Read a JSON-RPC message from stdin."""
        try:
            line = await asyncio.get_event_loop().run_in_executor(
                None, sys.stdin.readline
            )
            if not line:
                return None
            return json.loads(line.strip())
        except json.JSONDecodeError as e:
            logger.error(f"Failed to parse JSON: {e}")
            return None
        except Exception as e:
            logger.error(f"Error reading message: {e}")
            return None
    
    def _write_message(self, message: dict) -> None:
        """Write a JSON-RPC message to stdout."""
        sys.stdout.write(json.dumps(message) + "\n")
        sys.stdout.flush()
    
    def _create_response(self, id: Any, result: Any) -> dict:
        """Create a JSON-RPC response."""
        return {
            "jsonrpc": "2.0",
            "id": id,
            "result": result
        }
    
    def _create_error(self, id: Any, code: int, message: str) -> dict:
        """Create a JSON-RPC error response."""
        return {
            "jsonrpc": "2.0",
            "id": id,
            "error": {
                "code": code,
                "message": message
            }
        }
    
    async def _handle_request(self, request: dict) -> dict:
        """Handle an incoming MCP request."""
        method = request.get("method", "")
        params = request.get("params", {})
        request_id = request.get("id")
        
        # Initialize
        if method == "initialize":
            return self._create_response(request_id, {
                "protocolVersion": self.PROTOCOL_VERSION,
                "capabilities": {
                    "tools": {}
                },
                "serverInfo": {
                    "name": self.SERVER_NAME,
                    "version": self.SERVER_VERSION
                }
            })
        
        # List tools
        elif method == "tools/list":
            tools = [
                {
                    "name": tool.name,
                    "description": tool.description,
                    "inputSchema": tool.input_schema
                }
                for tool in self._tools.values()
            ]
            return self._create_response(request_id, {"tools": tools})
        
        # Call tool
        elif method == "tools/call":
            tool_name = params.get("name")
            arguments = params.get("arguments", {})
            
            if tool_name not in self._tools:
                return self._create_error(
                    request_id, -32601, f"Unknown tool: {tool_name}"
                )
            
            try:
                result = await self._tools[tool_name].handler(arguments)
                return self._create_response(request_id, {
                    "content": [
                        {
                            "type": "text",
                            "text": json.dumps(result, indent=2)
                        }
                    ]
                })
            except Exception as e:
                logger.exception(f"Error executing tool {tool_name}")
                return self._create_error(
                    request_id, -32603, f"Tool execution error: {str(e)}"
                )
        
        # Ping
        elif method == "ping":
            return self._create_response(request_id, {})
        
        # Notifications (no response needed)
        elif method == "notifications/initialized":
            return None
        
        # Unknown method
        else:
            return self._create_error(
                request_id, -32601, f"Method not found: {method}"
            )
    
    async def start(self) -> None:
        """Start the MCP server (stdio mode)."""
        logger.info(f"Starting {self.SERVER_NAME} v{self.SERVER_VERSION}")
        self._running = True
        
        while self._running:
            message = await self._read_message()
            if message is None:
                break
            
            response = await self._handle_request(message)
            if response is not None:
                self._write_message(response)
    
    def stop(self) -> None:
        """Stop the MCP server."""
        self._running = False


# =============================================================================
# CLI Entry Point
# =============================================================================

def main():
    """Command-line entry point for MCP server."""
    import argparse
    
    parser = argparse.ArgumentParser(
        description="QuinkGL MCP Server for gossip learning monitoring"
    )
    parser.add_argument(
        "--log-level",
        choices=["DEBUG", "INFO", "WARNING", "ERROR"],
        default="INFO",
        help="Logging level"
    )
    
    args = parser.parse_args()
    
    # Configure logging (to stderr to not interfere with MCP protocol)
    logging.basicConfig(
        level=getattr(logging, args.log_level),
        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
        stream=sys.stderr
    )
    
    # Create collector and server
    collector = MetricsCollector()
    server = MCPServer(collector)
    
    # Run
    try:
        asyncio.run(server.start())
    except KeyboardInterrupt:
        logger.info("Server stopped by user")


if __name__ == "__main__":
    main()
