from typing import Dict, Any
import aiohttp
from .base import Tool

class GoogleSearchTool(Tool):
    def __init__(self, api_key: str):
        super().__init__(
            name="google_search",
            description="Search the web using Google Custom Search API"
        )
        self.api_key = api_key
        self.cx = "your_custom_search_engine_id"  # Set your custom search engine ID
        
    async def execute(self, query: str) -> Dict[str, Any]:
        """Execute a Google search query"""
        url = "https://www.googleapis.com/customsearch/v1"
        params = {
            "key": self.api_key,
            "cx": self.cx,
            "q": query
        }
        
        async with aiohttp.ClientSession() as session:
            async with session.get(url, params=params) as response:
                data = await response.json()
                
                if "items" in data:
                    results = [{
                        "title": item.get("title"),
                        "link": item.get("link"),
                        "snippet": item.get("snippet")
                    } for item in data["items"][:5]]
                    
                    return {
                        "success": True,
                        "results": results
                    }
                    
                return {
                    "success": False,
                    "error": "No results found"
                }

class ExaSearchTool(Tool):
    def __init__(self, api_key: str):
        super().__init__(
            name="exa_search",
            description="Search the web using Exa API"
        )
        self.api_key = api_key
        
    async def execute(self, query: str) -> Dict[str, Any]:
        """Execute an Exa search query"""
        url = "https://api.exa.ai/search"
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }
        
        async with aiohttp.ClientSession() as session:
            async with session.post(
                url,
                json={"query": query},
                headers=headers
            ) as response:
                data = await response.json()
                
                if "results" in data:
                    return {
                        "success": True,
                        "results": data["results"]
                    }
                    
                return {
                    "success": False,
                    "error": "No results found"
                }
