import os
import gradio as gr
import tempfile
from langchain.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import FAISS
from langchain.chat_models import ChatOpenAI
from langchain.chains import ConversationalRetrievalChain

# Set your OpenAI API key
os.environ["OPENAI_API_KEY"] = "your-api-key-here"

class RAGChatbot:
    def __init__(self):
        self.vector_store = None
        self.qa_chain = None

    def process_pdf(self, pdf_file):
        # Save the uploaded file temporarily
        with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as temp_file:
            temp_file.write(pdf_file.read())
            temp_path = temp_file.name

        # Load the PDF
        loader = PyPDFLoader(temp_path)
        documents = loader.load()

        # Split the documents into chunks
        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
        splits = text_splitter.split_documents(documents)

        # Create embeddings and vector store
        embeddings = HuggingFaceEmbeddings()
        self.vector_store = FAISS.from_documents(splits, embeddings)

        # Create the conversational chain
        llm = ChatOpenAI(temperature=0, model_name="gpt-3.5-turbo")
        self.qa_chain = ConversationalRetrievalChain.from_llm(
            llm, retriever=self.vector_store.as_retriever(), return_source_documents=True
        )

        # Clean up the temporary file
        os.unlink(temp_path)

        return "PDF processed successfully. You can now ask questions about the document."

    def ask_question(self, question, history):
        if not self.qa_chain:
            return "Please upload a PDF document first.", []

        # Prepare the chat history
        chat_history = [(h[0], h[1]) for h in history]

        # Get the answer from the QA chain
        result = self.qa_chain({"question": question, "chat_history": chat_history})
        answer = result["answer"]
        source_docs = result["source_documents"]

        # Extract page numbers and content for highlighting
        highlights = []
        for doc in source_docs:
            page_number = doc.metadata.get("page", 0) + 1  # PyPDFLoader uses 0-based indexing
            content = doc.page_content
            highlights.append({"page": page_number, "content": content})

        return answer, highlights

chatbot = RAGChatbot()

def user(message, history):
    return "", history + [[message, None]]

def bot(history):
    question = history[-1][0]
    answer, highlights = chatbot.ask_question(question, history[:-1])
    history[-1][1] = answer
    return history, highlights

with gr.Blocks() as demo:
    gr.Markdown("# RAG Chatbot for PDF Documents")
    
    with gr.Row():
        with gr.Column(scale=1):
            chatbot = gr.Chatbot(height=600)
            msg = gr.Textbox(label="Ask a question about the PDF")
            clear = gr.Button("Clear")
        
        with gr.Column(scale=1):
            pdf_display = gr.HTML('<div id="pdf-viewer" style="height: 600px;"></div>')
    
    file_output = gr.Textbox(label="PDF Processing Status")
    upload_button = gr.UploadButton("Upload PDF", file_types=[".pdf"])
    
    highlights = gr.JSON(visible=False)
    
    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
        bot, chatbot, [chatbot, highlights]
    )
    clear.click(lambda: None, None, chatbot, queue=False)
    upload_button.upload(chatbot.process_pdf, upload_button, file_output)

    gr.HTML("""
    <script src="https://cdnjs.cloudflare.com/ajax/libs/pdf.js/2.9.359/pdf.min.js"></script>
    <script>
    let pdfDoc = null;
    let pageNum = 1;
    let pageRendering = false;
    let pageNumPending = null;
    let scale = 1.5;
    let canvas = null;
    let ctx = null;

    function renderPage(num) {
        pageRendering = true;
        pdfDoc.getPage(num).then(function(page) {
            let viewport = page.getViewport({scale: scale});
            canvas.height = viewport.height;
            canvas.width = viewport.width;
            let renderContext = {
                canvasContext: ctx,
                viewport: viewport
            };
            let renderTask = page.render(renderContext);
            renderTask.promise.then(function() {
                pageRendering = false;
                if (pageNumPending !== null) {
                    renderPage(pageNumPending);
                    pageNumPending = null;
                }
            });
        });
    }

    function queueRenderPage(num) {
        if (pageRendering) {
            pageNumPending = num;
        } else {
            renderPage(num);
        }
    }

    function highlightText(page, text) {
        pdfDoc.getPage(page).then(function(pdfPage) {
            pdfPage.getTextContent().then(function(textContent) {
                let textItems = textContent.items;
                let viewport = pdfPage.getViewport({scale: scale});
                ctx.save();
                ctx.fillStyle = "yellow";
                for (let i = 0; i < textItems.length; i++) {
                    if (textItems[i].str.includes(text)) {
                        let tx = textItems[i].transform;
                        let pos = viewport.convertToViewportPoint(tx[4], tx[5]);
                        ctx.fillRect(pos[0], pos[1], textItems[i].width * scale, textItems[i].height * scale);
                    }
                }
                ctx.restore();
                renderPage(page);
            });
        });
    }

    gradio_config.components.forEach(component => {
        if (component.props.label === "Upload PDF") {
            component.props.onChange = (files) => {
                if (files.length > 0) {
                    let fileReader = new FileReader();
                    fileReader.onload = function() {
                        let typedarray = new Uint8Array(this.result);
                        pdfjsLib.getDocument({data: typedarray}).promise.then(function(pdf) {
                            pdfDoc = pdf;
                            canvas = document.createElement('canvas');
                            ctx = canvas.getContext('2d');
                            document.getElementById('pdf-viewer').innerHTML = '';
                            document.getElementById('pdf-viewer').appendChild(canvas);
                            renderPage(pageNum);
                        });
                    };
                    fileReader.readAsArrayBuffer(files[0]);
                }
            };
        } else if (component.props.elem_id === "highlights") {
            component.props.onChange = (highlights) => {
                if (highlights && highlights.length > 0) {
                    highlights.forEach(highlight => {
                        highlightText(highlight.page, highlight.content);
                    });
                }
            };
        }
    });
    </script>
    """)

demo.launch()