import pandas as pd
from langchain_community.utilities import GoogleSerperAPIWrapper


class GoogleSerperSearchAgent:
    def __init__(self, sources_path: str) -> None:
        with open(sources_path) as f:
            self._sources = [line.strip() for line in f.readlines()]

        self._api = GoogleSerperAPIWrapper()

    def run(self, query: str, limit_per_source: int = 3) -> pd.DataFrame:
        results = []
        for site in self._sources:
            results.extend(self._search(query, site, limit_per_source))

        df = pd.DataFrame(results)

        return df

    def _search(self, query: str, site: str | None = None, limit: int = 3) -> list[dict[str, str]]:
        if site is not None:
            query += f" site:{site}"

        self._api.k = limit

        results = self._api.results(query)["organic"]
        results = [{"title": result["title"], "url": result["link"], "date": result.get("date", "")} for result in results]

        return results