import asyncio
from meta_researcher.tool.tools.search_engine.log import logger
import httpx
import aiohttp
from asyncache import cached as acached
from cachetools import TTLCache
from urllib.parse import urlparse

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

class GoogleBingOptionalSearch(BaseSearch):
    def __init__(self, args: SearchConfig):
        super().__init__(args)
        self.topk = 10
        self.name = "GoogleBingOptionalSearch"
        self.context_size = 30000
        self.base_url = args.base_url
        self.search_engine = args.optional_search_engine
        self.method = args.search_method
        self.return_full_search_content = True
        self.return_chat_prompt = False
        self.return_search_results_prompt = False
        self.return_search_results = True
        self.timeout = 10

    @acached(cache=TTLCache(maxsize=100, ttl=600))
    async def search(self, query: str) -> list[dict]:
        headers = {"Content-Type": "application/json"}
        payload = {
            "query": query,
            "context_size": self.context_size,
            "topk": self.topk,
            "return_full_search_content": self.return_full_search_content,
            "return_chat_prompt": self.return_chat_prompt,
            "return_search_results_prompt": self.return_search_results_prompt,
            "return_search_results": self.return_search_results,
            "search_engine": self.search_engine,
            "method": self.method,
            "timeout": self.timeout
        }
        for attempt in range(self.args.max_retry):
            try:
                logger.info(f"Begin search for {query} with GoogleBingOptionalSearch")
                async with httpx.AsyncClient(timeout=30.0) as client:
                    response = await client.post(
                        self.base_url + "/generate_prompt",
                        headers=headers,
                        json=payload,
                    )
                    response.raise_for_status()
                    result = response.json()
                return result
            except Exception as e:
                logger.warning(str(e))
                await asyncio.sleep(1)
        raise Exception('Failed to get search results from Bing Search after retries.')