from typing import List
from langchain_community.tools.ddg_search import DuckDuckGoSearchRun
from langchain.agents import AgentType, initialize_agent
from langchain.tools.base import BaseTool
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_core.messages import SystemMessage
from langchain.memory import ConversationBufferMemory

class GeminiAgent:
    def __init__(self, api_key: str, model_name: str = "gemini-2.0-flash"):
        # Suppress warnings
        import warnings
        warnings.filterwarnings("ignore", category=UserWarning)
        warnings.filterwarnings("ignore", category=DeprecationWarning)
        warnings.filterwarnings("ignore", message=".*will be deprecated.*")
        warnings.filterwarnings("ignore", "LangChain.*")
        
        self.api_key = api_key
        self.model_name = model_name
        self.agent = self._setup_agent()

    def _setup_agent(self):
        # Initialize model with system message
        model = ChatGoogleGenerativeAI(
            model=self.model_name,
            google_api_key=self.api_key,
            temperature=0,  # Lower temperature for faster, more focused responses
            max_output_tokens=200,  # Limit response length
            convert_system_message_to_human=True,  # Faster processing of system message
            stream=True,  # Enable streaming for faster initial response
            system_message=SystemMessage(content="You are a concise AI assistant. Provide a short and accurate answer. Preferable answer should be in one word or line. Unless if query asked expects an elaborate answer.")
        )

        # Setup tools
        tools: List[BaseTool] = [DuckDuckGoSearchRun()]

        # Setup memory
        memory = ConversationBufferMemory(
            memory_key="chat_history",
            return_messages=True
        )

        # Create and return agent
        return initialize_agent(
            tools,
            model,
            agent=AgentType.CHAT_CONVERSATIONAL_REACT_DESCRIPTION,
            memory=memory,
            verbose=False,
            handle_parsing_errors=True
        )

    def run(self, query: str) -> str:
        try:
            result = self.agent.invoke({"input": query})
            return result["output"]
        except Exception as e:
            return f"Error: {e}"

    def run_interactive(self):
        print("AI Assistant Ready! (Type 'exit' to quit)")
        
        while True:
            query = input("You: ").strip()
            if query.lower() == 'exit':
                print("Goodbye!")
                break
            
            print("Assistant:", self.run(query))