import os
import pymssql
import pandas as pd



from langchain.vectorstores import Chroma
from langchain.embeddings import OpenAIEmbeddings
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.llms import OpenAI
from langchain.chat_models import ChatOpenAI
from langchain.chains import RetrievalQA
from langchain.document_loaders import TextLoader
from langchain.document_loaders import DirectoryLoader
from langchain.document_loaders import CSVLoader
from langchain.memory import ConversationBufferMemory

def Loading():
    return "데이터 로딩 중..."

def LoadData(openai_key):
    
    if openai_key is not None:
        
        persist_directory = 'realdb_LLM'

        embedding = OpenAIEmbeddings()

        vectordb = Chroma(
            persist_directory=persist_directory, 
            embedding_function=embedding
        )

        global retriever
        retriever = vectordb.as_retriever(search_kwargs={"k": 1})
        
        return "준비 완료"
    else:
        return "사용하시는 API Key를 입력하여 주시기 바랍니다."

def process_llm_response(llm_response):
    print(llm_response['result'])
    print('\n\nSources:')
    for source in llm_response["source_documents"]:
        print(source.metadata['source'])
        
        
# 챗봇의 답변을 처리하는 함수
def respond(message, chat_history, temperature):
    try:
        qa_chain = RetrievalQA.from_chain_type(
            llm=OpenAI(temperature=0.4),
            # llm=OpenAI(temperature=0.4), 
            # llm=ChatOpenAI(temperature=0),
            chain_type="stuff", 
            retriever=retriever
        )

        result = qa_chain(message)
        
        bot_message = result['result']

        # bot_message += '\n\n' + ' [출처]'

        # # 답변의 출처를 표기
        # for i, doc in enumerate(result['source_documents']):
        #     bot_message += str(i+1) + '. ' + doc.metadata['source'] + ' '

        # 채팅 기록에 사용자의 메시지와 봇의 응답을 추가.
        chat_history.append((message, bot_message))
        
        return "", chat_history
    except:
        chat_history.append(("", "API Key 입력 요망"))
        
        return " ", chat_history

        # return "", chat_history

    
import gradio as gr

# 챗봇 설명
title = """
<div style="text-align: center; max-width: 500px; margin: 0 auto;">
    <div>
        <h1>Pretraining Chatbot V2 Real</h1>
    </div>
    <p style="margin-bottom: 10px; font-size: 94%">
        OpenAI LLM를 이용한 Chatbot (Similarity)
    </p>
</div>
"""

# 꾸미기
css="""
#col-container {max-width: 700px; margin-left: auto; margin-right: auto;}
"""
with gr.Blocks(css=css) as UnivChatbot:
    with gr.Column(elem_id="col-container"):
        gr.HTML(title)

        with gr.Row():
            with gr.Column(scale=3):
                openai_key = gr.Textbox(label="You OpenAI API key", type="password", placeholder="OpenAI Key Type", elem_id="InputKey", show_label=False, container=False)
            with gr.Column(scale=1):
                langchain_status = gr.Textbox(placeholder="Status", interactive=False, show_label=False, container=False)
            with gr.Column(scale=1):
                chk_key = gr.Button("확인", variant="primary")
                
        chatbot = gr.Chatbot(label="대학 챗봇시스템(OpenAI LLM)", elem_id="chatbot") # 상단 좌측 

        with gr.Row():
            with gr.Column(scale=9):
                msg = gr.Textbox(label="입력", placeholder="궁금하신 내역을 입력하여 주세요.", elem_id="InputQuery", show_label=False, container=False)
            
        with gr.Row():
            with gr.Column(scale=1):
                submit = gr.Button("전송", variant="primary")
            with gr.Column(scale=1):
                clear = gr.Button("초기화", variant="stop")

    #chk_key.click(Loading, None, langchain_status, queue=False)  
    chk_key.click(LoadData, openai_key, outputs=[langchain_status], queue=False)
    
    # 사용자의 입력을 제출(submit)하면 respond 함수가 호출.
    msg.submit(respond, [msg, chatbot], [msg, chatbot])

    submit.click(respond, [msg, chatbot], [msg, chatbot])

    # '초기화' 버튼을 클릭하면 채팅 기록을 초기화.
    clear.click(lambda: None, None, chatbot, queue=False)

  
UnivChatbot.launch()