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, ToolCodeExecution, GenerateContentConfig, Content, Part
from google import genai

class CodeExecutionInput(BaseModel):
    task: str = Field(description="The task to be executed")

class CodeExecutionTool(BaseTool):
    name: str = "code_execution"
    description: str = "Useful for generate and execute python code for a given task"
    args_schema: Optional[ArgsSchema] = CodeExecutionInput
    client: Any = None
    model_id: str = "gemini-2.0-flash"
    return_direct: bool = False
    code_execution_tool: Any = None

    def __init__(self, api_key: Optional[str] = None, **kwargs):
        super().__init__(**kwargs)
        self.client = genai.Client(api_key=api_key)
        self.code_execution_tool = Tool(code_execution=ToolCodeExecution())

    def _run(self, task: str, run_manager: Optional[CallbackManagerForToolRun]=None) -> str:
        """Generate and execute python code for a given task"""
        try:
            response = self.client.models.generate_content(
                model=self.model_id,
                contents=f"Task: {task} "
                "Plan on how to solve the task. Generate and run code to solve the task.",
                config=GenerateContentConfig(
                    tools=[self.code_execution_tool],
                    response_modalities=["TEXT"]
                )
            )

            result = ""
            for part in response.candidates[0].content.parts:
                if part.executable_code is not None:
                    result += part.executable_code.code
                if part.code_execution_result is not None:
                    result += "\n\n" + part.code_execution_result.output

            return result
        except Exception as e:
            return f"Error performing code execution: {str(e)}"

    async def _arun(self, task: str, run_manager: Optional[CallbackManagerForToolRun]=None) -> str:
        """Run the code execution tool asynchronously."""
        return self._run(task, 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
    code_execution = CodeExecutionTool()

    # Run a search query
    result = code_execution.invoke({"task": "What is the sum of the first 50 prime numbers?"})
    print(result)