import asyncio
from typing import List, Optional
from haystack import component, default_from_dict, default_to_dict
from haystack.dataclasses import Document
from brave_search_python_client import BraveSearch, WebSearchRequest

@component
class BraveWebSearch:
    def __init__(self, api_key: str, top_k: int = 5, country: Optional[str] = None):
        self.api_key = api_key
        self.top_k = top_k
        self.country = country

    @component.output_types(documents=List[Document])
    def run(self, query: str):
        return asyncio.run(self._async_run(query))

    async def _async_run(self, query: str):
        bs = BraveSearch(api_key=self.api_key)
        req = WebSearchRequest(q=query, count=self.top_k, country=self.country)
        response = await bs.web(req)
        results = response.web.results if response.web else []
        return {"documents": [
            Document(content=r.description or "", meta={"title": r.title, "url": str(r.url)})
            for r in results
        ]}

    def to_dict(self):
        return default_to_dict(self, api_key=self.api_key, top_k=self.top_k, country=self.country)

    @classmethod
    def from_dict(cls, data):
        return default_from_dict(cls, data)
