import os
import gradio as gr
import openai
from langchain import hub
from langchain_community.document_loaders import PyPDFLoader
from langchain_community.vectorstores import Chroma
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter


data_root = './data/pdf/'
pdf_paths = [data_root+path for path in os.listdir(data_root)]

loaders = [PyPDFLoader(path) for path in pdf_paths]

docs = []
for loader in loaders:
    docs.extend(
        loader.load()[0:] # skip first page
    )

chunk_size = 1000
chunk_overlap = 200

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

splits = text_splitter.split_documents(docs)

vectorstore = Chroma.from_documents(documents=splits, embedding=OpenAIEmbeddings())
retriever = vectorstore.as_retriever()
prompt = hub.pull("rlm/rag-prompt")

model_name = 'gpt-3.5-turbo-0125'
# model_name = 'gpt-4-1106-preview'
model_name = 'gpt-4-0125-preview'
llm = ChatOpenAI(model_name=model_name, temperature=0)

def format_docs(docs):
    return '\n\n'.join(doc.page_content for doc in docs)

rag_chain = (
    {"context": retriever | format_docs, "question": RunnablePassthrough()}
    | prompt
    | llm
    | StrOutputParser()
)

def predict(query):
    return rag_chain.invoke(query)

examples = [
    "هل هناك غرامة للتخلف عن سداد ضريبة القيمة المضافة؟",
    "ما هي ضريبة القيمة المضافة؟",
    "ما الواجب على الخاضغين لضريبة القيمة المضافة؟",
    "من هو الشخص الخاضغ لضريبة القيمة المضافة؟",
    "متى يجب على الشخص التسجيل لضريبة القيمة المضافة؟",
    "أريد بيع منزل, هل يخضع ذلك لضريبة القيمة المضافة؟"
]

textbox = gr.Textbox(label="اكتب سؤالك هنا", placeholder="", lines=4)
iface = gr.Interface(fn=predict, inputs=textbox, outputs="text", examples=examples)
iface.launch(share=True)