import pandas as pd

from .dataset_semantic_search_agent import DatasetSemanticSearchAgent
from .google_serper_search_agent import GoogleSerperSearchAgent


class ResearchAgent:
    def __init__(self, web_sources_path: str = "sources.txt", local_sources_path: str = "article_summaries") -> None:
        self._google_serper_search_agent = GoogleSerperSearchAgent(web_sources_path)
        self._dataset_semantic_search_agent = DatasetSemanticSearchAgent(local_sources_path)

    def run(self, themes: list[dict[str, str]]) -> tuple[dict[str, dict[str, pd.DataFrame]], dict[str, float]]:
        sources = {}
        metrics = {"cost": 0}

        for theme in themes:
            web_sources, metrics_a = self._google_serper_search_agent.run(theme["query"])
            dataset_sources, metrics_b = self._dataset_semantic_search_agent.run(theme["query"])

            sources[theme["title"]] = {"web_sources": web_sources, "dataset_sources": dataset_sources}
            metrics["cost"] += metrics_a["cost"] + metrics_b["cost"]

        return sources, metrics