# https://qiita.com/nekoniii3/items/5acf764af65212d9f04f

import gradio as gr
import random
import time

import os

from langchain_community.document_loaders import PyMuPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
# from langchain_community.chat_models import ChatOpenAI
from langchain_openai import ChatOpenAI
from langchain_community.vectorstores import Chroma
from langchain.chains import RetrievalQA
from langchain_community.embeddings import OpenAIEmbeddings



os.environ["TOKENIZERS_PARALLELISM"] = "false"
os.environ["OPENAI_API_KEY"] = "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"

file_name1 = 'ALV2_ALV3DTU操作マニュアルDTU-V3SET01.pdf'
file_name2 = 'ALV3PCサーバ_ソフトウェア操作マニュアル_画像ファイル名付.pdf'
file_name3 = '美和ロック総合カタログ第31版_前半.pdf'
file_name4 = '美和ロック総合カタログ第31版_後半.pdf'

loader1 = PyMuPDFLoader(file_name1)
loader2 = PyMuPDFLoader(file_name2)
loader3 = PyMuPDFLoader(file_name3)
loader4 = PyMuPDFLoader(file_name4)

documents1 = loader1.load()
documents2 = loader2.load()
documents3 = loader3.load()
documents4 = loader4.load()

text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=0)

texts1 = text_splitter.split_documents(documents1)
texts2 = text_splitter.split_documents(documents2)
texts3 = text_splitter.split_documents(documents3)
texts4 = text_splitter.split_documents(documents4)
texts = texts1 + texts2 + texts3 + texts4

embeddings = OpenAIEmbeddings(model="text-embedding-ada-002") 
vectordb = Chroma.from_documents(texts, embeddings)
llm = ChatOpenAI(model_name="gpt-3.5-turbo-16k", temperature=0.05)

qa = RetrievalQA.from_chain_type(
    llm=llm, 
    chain_type="stuff", 
    retriever=vectordb.as_retriever(),
    return_source_documents=True)

import shutil
def save_image_filepath(filepath: str):
    print(filepath)
    # イメージを保存
    _, file_extension = os.path.splitext(filepath)
    shutil.copy(filepath, './filepath{}'.format(file_extension))
    pass

with gr.Blocks() as demo:
    chatbot = gr.Chatbot()
    # with gr.Row():
    #     with gr.Column():
    #         image_input_filepath = gr.Image(type='filepath')
            # image_button_filepath = gr.Button("filepath")

    # image_button_filepath.click(save_image_filepath, inputs=image_input_filepath)

    msg = gr.Textbox()

    def user(user_message, history):
        reply2 = qa(user_message)
        reply=reply2['result']

        for sd in reply2["source_documents"]:
            page_content = str(sd.page_content)
            source = str(sd.metadata["source"])
            page = str(sd.metadata["page"]+1).zfill(3)
            print("PDF：" + source)
            print("ページ：" + page)
        
        reply = reply + '<a href="https://dcs.mediapress-net.com/iportal/cv.do?c=20958580000&pg=108&v=MIW10001&d=LINK_MIW">link</a>'

        return "", history + [[user_message, reply]]

    def bot(history):
        yield history
        # save_image_filepath("./IMG_yosuke2.jpg")

    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=True).then(
        bot, chatbot, chatbot
    )
    
demo.queue()
demo.launch(share=True)