from typing import Any, Optional
from langchain_core.callbacks import CallbackManagerForToolRun
from langchain_core.tools import BaseTool
from langchain_core.tools.base import ArgsSchema
from pydantic import BaseModel, Field
from google.genai.types import Tool, GoogleSearch, GenerateContentConfig, Content, Part
from google import genai

class GoogleSearchInput(BaseModel):
    query: str = Field(description="The query to search for")

class GoogleSearchTool(BaseTool):
    name: str = "google_search"
    description: str = "Useful for searching Google to find information about current events, data, or public information."
    args_schema: Optional[ArgsSchema] = GoogleSearchInput
    client: Any = None
    model_id: str = "gemini-2.0-flash"
    return_direct: bool = False
    google_search_tool: Any = None

    def __init__(self, api_key: Optional[str] = None, **kwargs):
        super().__init__(**kwargs)
        self.client = genai.Client(api_key=api_key)
        self.google_search_tool = Tool(google_search=GoogleSearch())

    def _run(self, query: str, run_manager: Optional[CallbackManagerForToolRun]=None) -> str:
        """Run the Google search with the query."""
        try:
            response = self.client.models.generate_content(
                model=self.model_id,
                contents=Content(
                    parts=[Part.from_text(text=f"Search for information about: {query}")],
                    role="user"
                ),
                config=GenerateContentConfig(
                    tools=[self.google_search_tool],
                    response_modalities=["TEXT"]
                )
            )

            # Extract search results
            return response.text
        except Exception as e:
            return f"Error performing Google search: {str(e)}"

    async def _arun(self, query: str, run_manager: Optional[CallbackManagerForToolRun]=None) -> str:
        """Run the Google search with the query asynchronously."""
        return self._run(query, run_manager=run_manager.get_sync())

# Example usage
if __name__ == "__main__":
    from dotenv import load_dotenv

    # Load API key from environment variables
    load_dotenv()

    # Create the Google Search runnable
    google_search = GoogleSearchTool()

    # Run a search query
    result = google_search.invoke({"query": "What is the latest news about AI?"})
    print(result)