from dotenv import load_dotenv

# langchain libraries
from langchain.chat_models import ChatOpenAI
from langchain.chains import RetrievalQAWithSourcesChain
from langchain.prompts.chat import (
    ChatPromptTemplate,
    SystemMessagePromptTemplate,
    HumanMessagePromptTemplate,
)
import openai
import os

from vector_store import get_or_create_vector_store
from pinecone_vector_store import get_pinecone_store


load_dotenv()
openai.api_key = os.getenv("OPENAI_API_KEY")


def generate_prompt_templates():
    system_template = """Use the following pieces of context to answer the users question as long as possible.  Provided context is regulations of subway company in busan. what users want to query is exact details or inferable information from regulations.
Given the following summaries of a long document and a question, create a final answer with references to the document.
If you don't know the answer, just say that "죄송합니다. 해당질문에 대한 답을 찾지 못했습니다.😿\\n 사규에 포함된 내용을 구체적으로 범위를 좁혀서 질문해주세요.🤖\\n 원하는 답을 얻지 못했을 경우 어떤 규정을 대상으로 질문하는지 명시해서 다시 질문해 보세요! ", don't try to make up an answer. + '\n'
----------------
{summaries}

You MUST answer in Korean:"""

    messages = [
        SystemMessagePromptTemplate.from_template(system_template),
        HumanMessagePromptTemplate.from_template("{question}")
    ]

    prompt = ChatPromptTemplate.from_messages(messages)
    return prompt


def generate_chain(pincone=True):
    if pincone:
        vector_store = get_pinecone_store()
    else:
        vector_store = get_or_create_vector_store()

    retriever = vector_store.as_retriever(search_kwargs={"k": 4})

    llm = ChatOpenAI(model_name="gpt-3.5-turbo-16k", temperature=0.3)
    chain_type_kwargs = {"prompt": generate_prompt_templates()}
    chain = RetrievalQAWithSourcesChain.from_chain_type(
        llm=llm,
        chain_type="stuff",
        retriever=retriever,
        return_source_documents=True,
        chain_type_kwargs=chain_type_kwargs)
    return chain