# Ingest script
# File: scripts/ingest.py
import yaml
import sys
from orchestrator.client import MCPClient


def main():
    """
    Ingest papers for a given query into Chroma vector store via MCP.
    Usage: python ingest.py "your search query"
    """
    if len(sys.argv) < 2:
        print("Usage: python ingest.py <query>")
        sys.exit(1)
    query = sys.argv[1]

    cfg = yaml.safe_load(open("config.yaml"))
    web = MCPClient(cfg['mcp_servers']['web_search'])
    pubmed = MCPClient(cfg['mcp_servers']['pubmed'])
    chroma = MCPClient(cfg['mcp_servers']['chroma'])

    print(f"Ingesting papers for query: {query}")
    results = []
    try:
        results += web.call("web_search.search", {"q": query})
    except Exception as e:
        print("Web search failed:", e)
    try:
        results += pubmed.call("metatool.query", {"source": "PubMed", "q": query})
    except Exception as e:
        print("PubMed search failed:", e)

    for paper in results:
        paper_id = paper.get('id')
        text = paper.get('abstract', '')
        meta = {"title": paper.get('title'), "authors": ",".join(paper.get('authors', []))}
        chroma.call("chroma.insert", {"id": paper_id, "text": text, "metadata": meta})
    print("Ingestion complete!")

if __name__ == "__main__":
    main()