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

class Dispatcher:
    """
    Coordinates calls to MCP servers for search, notebook execution, and graph retrieval.
    """
    def __init__(self, config_path='config.yaml'):
        with open(config_path, 'r') as f:
            cfg = yaml.safe_load(f)
        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):
        """
        Aggregate results from web and PubMed MCP endpoints.
        """
        papers = []
        try:
            papers += self.web.call('web_search.search', {'q': query}) or []
        except Exception:
            pass
        try:
            papers += self.pubmed.call('metatool.query', {'source': 'PubMed', 'q': query}) or []
        except Exception:
            pass
        # dedupe
        unique = {p['id']: p for p in papers if 'id' in p}
        return list(unique.values())[:limit]

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

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