import asyncio
from meta_researcher.tool.tools.search_engine.log import logger

import httpx
from asyncache import cached as acached
from cachetools import TTLCache

from meta_researcher.tool.tools.search_engine.base_search import BaseSearch, SearchConfig

class GoogleSearch(BaseSearch):
    def __init__(self, args: SearchConfig):
        super().__init__(args)
        self.topk = 10
        self.name = "GoogleSearch"

    async def _search(self, query: str) -> dict:
        serper_search_url = "https://www.googleapis.com/customsearch/v1?"
        params = {
            "q": query,
            "key": self.args.google_search_key,
            "cx": "27f79a5f8b6c3413e",
            "cr": "countryCN",
            "gl": "cn",
            "lr": "lang_zh-CN",
            "num": self.topk,
        }
        headers = {"Content-Type": "application/json"}
        async with httpx.AsyncClient(timeout=self.args.timeout) as client:
            response = await client.get(serper_search_url, headers=headers, params=params)
            response.raise_for_status()
            return response.json()

    def _parse_response(self, response) -> list[dict]:
        results = []
        for item in response["items"]:
            results.append(
                {
                    "url": item.get("formattedUrl", item["link"]),
                    "summ": item.get("snippet", ""),
                    "title": item.get("title", ""),
                    "meta": {
                        "from": self.name,
                    },
                }
            )
        return self._filter_results(results)

    @acached(cache=TTLCache(maxsize=100, ttl=600))
    async def search(self, query: str) -> list[dict]:
        for attempt in range(self.args.max_retry):
            try:
                response = await self._search(query)
                return self._parse_response(response)
            except Exception as e:
                logger.warning(str(e))
                await asyncio.sleep(1)
        raise Exception("Failed to get search results from Google Serper Search after retries.")
