from tools.base_tool import BaseTool
from googleapiclient.discovery import build
import os

class GoogleSearchTool(BaseTool):
    name = "web_search"
    description = "Performs a Google web search for a query and returns the top results in markdown format."

    inputs = {
        "query": {
            "type": "string",
            "description": "Search query for Google."
        }
    }

    output_type = "string"
    skip_forward_signature_validation = True

    def __init__(
        self,
        api_key: str | None = None,
        search_engine_id: str | None = None,
        num_results: int = 10,
        **kwargs,
    ):
        api_key = api_key or os.getenv("GOOGLE_SEARCH_API_KEY")
        search_engine_id = search_engine_id or os.getenv("GOOGLE_SEARCH_ENGINE_ID")

        if not api_key or not search_engine_id:
            raise EnvironmentError("Missing Google Search API credentials.")

        self.search_service = build("customsearch", "v1", developerKey=api_key).cse()
        self.engine_id = search_engine_id
        self.num_results = num_results

        super().__init__(**kwargs)

    def _collect_params(self) -> dict:
        return {}

    def forward(self, query: str, *args, **kwargs) -> str:
        search_args = {
            "q": query,
            "cx": self.engine_id,
            "fields": "items(link,title,snippet)",
            "num": self.num_results,
        }
        search_args.update(self._collect_params(*args, **kwargs))

        response = self.search_service.list(**search_args).execute()
        items = response.get("items", [])

        if not items:
            return "No results found."

        results = [f"**{item['title']}**\n{item['link']}\n{item['snippet']}" for item in items]
        return "\n\n".join(results)

class GoogleSiteSearchTool(GoogleSearchTool):
    name = "site_search"
    description = "Searches a specific domain for a query using Google and returns results."

    inputs = {
        "query": {
            "type": "string",
            "description": "Search query."
        },
        "site": {
            "type": "string",
            "description": "Domain to restrict search to."
        },
    }

    def _collect_params(self, site: str) -> dict:
        return {
            "siteSearch": site,
            "siteSearchFilter": "i"
        }