import asyncio
import json
import sys

class CryptoMCPServer:
    def __init__(self):
        self.tools = {}

    def register_tool(self, name: str, handler: callable, input_schema: dict, output_schema: dict, streaming: bool = False):
        self.tools[name] = {
            "handler": handler,
            "input_schema": input_schema,
            "output_schema": output_schema,
            "streaming": streaming
        }

    async def run(self):
        print("Crypto MCP Server is running...", file=sys.stderr)

        while True:
            try:
                raw = await asyncio.get_event_loop().run_in_executor(
                    None, sys.stdin.readline
                )

                if not raw:
                    continue

                request = json.loads(raw.strip())

                method = request.get("method")
                params = request.get("params", {})
                request_id = request.get("id")

                # Handle initialize request
                if method == "initialize":
                    response = {
                        "jsonrpc": "2.0",
                        "id": request_id,
                        "result": {
                            "protocolVersion": "2024-11-05",
                            "serverInfo": {
                                "name": "Crypto MCP Server",
                                "version": "0.1.0"
                            },
                            "capabilities": {
                                "tools": {}
                            }
                        }
                    }
                    sys.stdout.write(json.dumps(response) + "\n")
                    sys.stdout.flush()

                # Handle initialized notification
                elif method == "notifications/initialized":
                    # Just acknowledge
                    pass

                # Handle tools/list request
                elif method == "tools/list":
                    tools_list = []
                    for name, tool_info in self.tools.items():
                        tools_list.append({
                            "name": name,
                            "description": f"Tool: {name}",
                            "inputSchema": tool_info["input_schema"]
                        })
                    response = {
                        "jsonrpc": "2.0",
                        "id": request_id,
                        "result": {
                            "tools": tools_list
                        }
                    }
                    sys.stdout.write(json.dumps(response) + "\n")
                    sys.stdout.flush()

                # Handle tool calls
                elif method == "tools/call":
                    tool_name = params.get("name")
                    tool_args = params.get("arguments", {})

                    if tool_name not in self.tools:
                        response = {
                            "jsonrpc": "2.0",
                            "id": request_id,
                            "error": {
                                "code": -32601,
                                "message": f"Tool '{tool_name}' not found"
                            }
                        }
                    else:
                        handler = self.tools[tool_name]["handler"]
                        streaming = self.tools[tool_name]["streaming"]

                        try:
                            if streaming:
                                # For streaming tools, we need to handle async generators
                                result = []
                                async for item in handler(**tool_args):
                                    result.append(item)
                                response = {
                                    "jsonrpc": "2.0",
                                    "id": request_id,
                                    "result": {
                                        "content": [
                                            {
                                                "type": "text",
                                                "text": json.dumps(result)
                                            }
                                        ]
                                    }
                                }
                            else:
                                result = await handler(**tool_args)
                                response = {
                                    "jsonrpc": "2.0",
                                    "id": request_id,
                                    "result": {
                                        "content": [
                                            {
                                                "type": "text",
                                                "text": json.dumps(result)
                                            }
                                        ]
                                    }
                                }
                        except Exception as exc:
                            response = {
                                "jsonrpc": "2.0",
                                "id": request_id,
                                "error": {
                                    "code": -32603,
                                    "message": str(exc)
                                }
                            }

                    sys.stdout.write(json.dumps(response) + "\n")
                    sys.stdout.flush()

                else:
                    # Unknown method
                    response = {
                        "jsonrpc": "2.0",
                        "id": request_id,
                        "error": {
                            "code": -32601,
                            "message": f"Method not found: {method}"
                        }
                    }
                    sys.stdout.write(json.dumps(response) + "\n")
                    sys.stdout.flush()

            except json.JSONDecodeError:
                print("Server error: Invalid JSON", file=sys.stderr)
            except Exception as e:
                print(f"Server error: {e}", file=sys.stderr)