from langchain.agents import load_tools
from langchain.agents import ZeroShotAgent, Tool, AgentExecutor
from langchain.memory import ConversationBufferMemory
from langchain.agents import initialize_agent
from langchain.agents import AgentType
from langchain import LLMChain


class CompletionMeta(type):
    """
    The meta class for completion interface
    """

tools = None

class Completion():
    def __init__(self):
        self.prefix = None
        self.suffix = None
        self.promptVars = None

    def complete(self, message: str) -> str:
        """
        Perform a text completion using the language model
        """
        pass

    def getModelName(self) -> str:
        """
        Return the model name
        """
        pass

    def getChain(self, llm):
        pass

    def getAgent(self, llm):
        # Load tools
        tools = load_tools(["serpapi", "llm-math"], llm=llm)

        # Build a prompt
        prompt = ZeroShotAgent.create_prompt(
                tools=tools,
                prefix=self.getPromptPrefix(),
                suffix=self.getPromptSuffix(),
                input_variables=self.getPromptVariables()
                )
        memory = ConversationBufferMemory(memory_key="chat_history")

        # Build the LLM Chain
        llm_chain = LLMChain(llm=llm, prompt=prompt)
        agent = ZeroShotAgent(llm_chain=llm_chain,
                              tools=tools,
                              verbose=True)
        self.agent_chain = AgentExecutor.from_agent_and_tools(agent=agent,
                                                         tools=tools,
                                                         verbose=True,
                                                         memory=memory)

        return self.agent_chain

    def getPromptVariables(self):
        if self.promptVars is None:
            self.promptVars = ["input", "chat_history", "agent_scratchpad"]
        return self.promptVars

    def getPromptPrefix(self):
        if self.prefix is None:
            self.prefix = """Have a conversation with a human, answering the following
            questions as best you can. You have access to the following tools:"""
        return self.prefix

    def getPromptSuffix(self):
        if self.suffix is None:
            self.suffix = """Begin!
            {chat_history}
            Question: {input}
            {agent_scratchpad}"""
        return self.suffix