from langchain.agents import Tool
from langchain.memory import ConversationBufferMemory ,ConversationBufferWindowMemory
from langchain.chat_models import ChatOpenAI
from langchain.agents import initialize_agent
from llama_index import GPTSimpleVectorIndex
import os
from langchain.schema import (
    SystemMessage
)



class ChatBot:
    def __init__(self, memory, agent_chain):
        self.memory = memory
        self.agent = agent_chain


def create_chatbot(model_name, seed_memory=None):
    # search = GoogleSearchAPIWrapper()
    # tools = [
    #     Tool(
    #         name="Current Search",
    #         func=search.run,
    #         description="useful for all question that asks about live events",
    #     ),
    #     Tool(
    #         name="Topic Search",
    #         func=search.run,
    #         description="useful for all question that are related to a particular topic, product, concept, or service",
    #     )
    # ]
    index = GPTSimpleVectorIndex.load_from_disk('martin.json')
    query_mode ="svm"

    tools = [
        Tool(
            name="GPT Index",
            func=lambda q: str(index.query(q,vector_store_query_mode=query_mode)),
            description="useful for when you want to answer questions about Martin Seligman and positive psychonogy related. The input to this tool should be a complete english sentence.",
            return_direct=True
        ),
    ]

    # messages = [
    #     SystemMessage(content="You are Martin Seligman. You use a tone that is warm and kind.")
    # ]
    #memory = ConversationBufferMemory(memory_key="chat_history")
    memory = seed_memory if seed_memory is not None else ConversationBufferWindowMemory( k=4 ,memory_key="chat_history")
    #memory = seed_memory if seed_memory is not None else ConversationBufferMemory(memory_key="chat_history")
    chat = ChatOpenAI(temperature=0, model_name=model_name)
    agent_chain = initialize_agent(tools, chat, agent="conversational-react-description", verbose=True, memory=memory)

    return ChatBot(memory, agent_chain)