"""
Script for the multilingual conversation based experimental AI agent
We used ChatGPT as the language model
You need openai api key to run this script
"""

import argparse
import os

from colossalqa.data_loader.document_loader import DocumentLoader
from colossalqa.data_loader.table_dataloader import TableLoader
from langchain import LLMChain, OpenAI
from langchain.agents import Tool, ZeroShotAgent
from langchain.agents.agent import AgentExecutor
from langchain.agents.agent_toolkits import create_retriever_tool
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.llms import OpenAI
from langchain.memory import ChatMessageHistory, ConversationBufferMemory
from langchain.memory.chat_memory import ChatMessageHistory
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.utilities import SQLDatabase
from langchain.vectorstores import Chroma
from langchain_experimental.sql import SQLDatabaseChain

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Experimental AI agent powered by ChatGPT")
    parser.add_argument("--open_ai_key_path", type=str, default=None, help="path to the plain text open_ai_key file")

    args = parser.parse_args()

    # Setup openai key
    # Set env var OPENAI_API_KEY or load from a file
    openai_key = open(args.open_ai_key_path).read()
    os.environ["OPENAI_API_KEY"] = openai_key

    # Load data served on sql
    print("Select files for constructing sql database")
    tools = []

    llm = OpenAI(temperature=0.0)

    while True:
        file = input("Select a file to load or press Enter to exit:")
        if file == "":
            break
        data_name = input("Enter a short description of the data:")

        table_loader = TableLoader(
            [[file, data_name.replace(" ", "_")]], sql_path=f"sqlite:///{data_name.replace(' ', '_')}.db"
        )
        sql_path = table_loader.get_sql_path()

        # Create sql database
        db = SQLDatabase.from_uri(sql_path)
        print(db.get_table_info())

        db_chain = SQLDatabaseChain.from_llm(llm, db, verbose=True)
        name = f"Query the SQL database regarding {data_name}"
        description = (
            f"useful for when you need to answer questions based on data stored on a SQL database regarding {data_name}"
        )
        tools.append(
            Tool(
                name=name,
                func=db_chain.run,
                description=description,
            )
        )
        print(f"Added sql dataset\n\tname={name}\n\tdescription:{description}")

    # VectorDB
    embedding = OpenAIEmbeddings()

    # Load data serve on sql
    print("Select files for constructing retriever")
    while True:
        file = input("Select a file to load or press Enter to exit:")
        if file == "":
            break
        data_name = input("Enter a short description of the data:")
        retriever_data = DocumentLoader([[file, data_name.replace(" ", "_")]]).all_data

        # Split
        text_splitter = RecursiveCharacterTextSplitter(chunk_size=200, chunk_overlap=0)
        splits = text_splitter.split_documents(retriever_data)

        # Create vector store
        vectordb = Chroma.from_documents(documents=splits, embedding=embedding)
        # Create retriever
        retriever = vectordb.as_retriever(
            search_type="similarity_score_threshold", search_kwargs={"score_threshold": 0.5, "k": 5}
        )
        # Add to tool chain
        name = f"Searches and returns documents regarding {data_name}."
        tools.append(create_retriever_tool(retriever, data_name, name))

    prefix = """Have a conversation with a human, answering the following questions as best you can. You have access to the following tools. If none of the tools can be used to answer the question. Do not share uncertain answer unless you think answering the question doesn't need any background information. In that case, try to answer the question directly."""
    suffix = """You are provided with the following background knowledge:
    Begin!"

    {chat_history}
    Question: {input}
    {agent_scratchpad}"""

    prompt = ZeroShotAgent.create_prompt(
        tools,
        prefix=prefix,
        suffix=suffix,
        input_variables=["input", "chat_history", "agent_scratchpad"],
    )

    memory = ConversationBufferMemory(memory_key="chat_history", chat_memory=ChatMessageHistory())

    llm_chain = LLMChain(llm=OpenAI(temperature=0.7), prompt=prompt)
    agent = ZeroShotAgent(llm_chain=llm_chain, tools=tools, verbose=True)
    agent_chain = AgentExecutor.from_agent_and_tools(agent=agent, tools=tools, verbose=True, memory=memory)

    while True:
        user_input = input("User: ")
        if " end " in user_input:
            print("Agent: Happy to chat with you ：)")
            break
        agent_response = agent_chain.run(user_input)
        print(f"Agent: {agent_response}")
    table_loader.sql_engine.dispose()