"""
MCP Server implementation for Medical CDM Tools.
This server exposes the TrendAnalyzer and FHIR tools to any MCP-compliant client.
"""
import sys
import json
import asyncio
from typing import Any, Dict, List
import src.tools.fhir_memory as fhir_tools
import src.tools.dietary_tools as dietary_tools
from langchain_core.tools import BaseTool
from src.agents.cdm_agents import TrendAnalyzer
from src.utils.logger import setup_logger

logger = setup_logger("MCPServer")

class MedicalMCPServer:
    def __init__(self):
        self.trend_analyzer = TrendAnalyzer()
        self.tools = {}
        for module in [fhir_tools, dietary_tools]:
            for name in dir(module):
                obj = getattr(module, name)
                if isinstance(obj, BaseTool):
                    self.tools[obj.name] = obj

    async def handle_request(self, request: Dict[str, Any]) -> Dict[str, Any]:
        method = request.get("method")
        params = request.get("params", {})
        req_id = request.get("id")
        logger.info(f"Received MCP request: method={method}, id={req_id}")

        try:
            if method == "list_tools":
                result = self.list_tools()
            elif method == "call_tool":
                result = await self.call_tool(params.get("name"), params.get("arguments", {}))
            else:
                return {"jsonrpc": "2.0", "error": {"code": -32601, "message": "Method not found"}, "id": req_id}
            
            return {"jsonrpc": "2.0", "result": result, "id": req_id}
        except Exception as e:
            return {"jsonrpc": "2.0", "error": {"code": -32603, "message": str(e)}, "id": req_id}

    def list_tools(self) -> List[Dict[str, Any]]:
        tool_list = [
            {
                "name": "analyze_health_trends",
                "description": "Analyze FHIR observation trends for a patient.",
                "inputSchema": {
                    "type": "object",
                    "properties": {
                        "patient_id": {"type": "string"}
                    },
                    "required": ["patient_id"]
                }
            }
        ]
        
        for name, tool_obj in self.tools.items():
            schema = {}
            if tool_obj.args_schema:
                try:
                    schema = tool_obj.args_schema.schema()
                except AttributeError:
                    schema = tool_obj.args_schema.model_json_schema()
            
            tool_list.append({
                "name": name,
                "description": tool_obj.description,
                "inputSchema": schema
            })
            
        return tool_list

    async def call_tool(self, name: str, args: Dict[str, Any]) -> Any:
        logger.info(f"Calling MCP tool: {name}")
        if name == "analyze_health_trends":
            return await self.trend_analyzer.analyze_trends(args["patient_id"])
        elif name in self.tools:
            return self.tools[name].invoke(args)
        else:
            raise ValueError(f"Unknown tool: {name}")

async def main():
    server = MedicalMCPServer()
    # Simple stdio loop for MCP
    while True:
        line = await asyncio.get_event_loop().run_in_executor(None, sys.stdin.readline)
        if not line:
            break
        try:
            request = json.loads(line)
            response = await server.handle_request(request)
            print(json.dumps(response), flush=True)
        except Exception as e:
            print(json.dumps({"error": str(e)}), flush=True)

if __name__ == "__main__":
    if len(sys.argv) > 1 and sys.argv[1] == "--serve":
        asyncio.run(main())
    else:
        print("Medical MCP Server. Use --serve to start in stdio mode.")