# File: orchestrator/dispatcher.py
import uuid
import yaml
from orchestrator.client import MCPClient


class Dispatcher:
    """
    Coordinates calls to various MCP servers for searching, code execution, and graph retrieval.
    """
    def __init__(self, config_path="config.yaml"):
        cfg = yaml.safe_load(open(config_path))
        servers = cfg.get("mcp_servers", {})
        self.web = MCPClient(servers.get("web_search"))
        self.pubmed = MCPClient(servers.get("pubmed"))
        self.chroma = MCPClient(servers.get("chroma"))
        self.runner = MCPClient(servers.get("python_run"))

    def search_papers(self, query: str, limit: int = 5):
        """
        Fan out search requests to web and PubMed MCP servers, aggregate and limit.
        """
        results = []
        try:
            results += self.web.call("web_search.search", {"q": query})
        except Exception:
            pass
        try:
            results += self.pubmed.call("metatool.query", {"source": "PubMed", "q": query})
        except Exception:
            pass
        # Deduplicate by id
        unique = {paper["id"]: paper for paper in results}
        return list(unique.values())[:limit]

    def get_notebook_cells(self, paper_id: str):
        """
        Retrieve code cells for reproducible example from the Python-run MCP server.
        """
        try:
            resp = self.runner.call("mcp-run-python.execute", {"paper_id": paper_id})
            return resp.get("cells", [])
        except Exception:
            return []

    def get_graph(self, paper_id: str):
        """
        Retrieve a knowledge graph representation for a paper from the Chroma MCP server.
        """
        try:
            return self.chroma.call("chroma.query_graph", {"id": paper_id})
        except Exception:
            return {"nodes": [], "edges": []}