import gradio as gr
import streamlit as st
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.text_splitter import CharacterTextSplitter
from langchain.vectorstores import Chroma
from langchain.chains import ConversationalRetrievalChain
from langchain.chat_models import ChatOpenAI
from langchain.document_loaders import PyPDFLoader
import os
import fitz
from PIL import Image

# Global variables
COUNT, N = 0, 0
chat_history = []
chain = None  # Initialize chain as None

# Function to set the OpenAI API key
def set_apikey(api_key):
    os.environ['OPENAI_API_KEY'] = api_key
    return disable_box  

# Function to enable the API key input box
def enable_api_box():
    return enable_box

# Function to add text to the chat history
def add_text(history, text):
    if not text:
        raise gr.Error('Enter text')
    history = history + [(text, '')]
    return history

# Function to process the PDF file and create a conversation chain
def process_file(file):
    global chain
    if 'OPENAI_API_KEY' not in os.environ:
        raise gr.Error('Upload your OpenAI API key')

    # Replace with your actual PDF processing logic
    loader = PyPDFLoader(file.name)
    documents = loader.load()
    embeddings = OpenAIEmbeddings()
    pdfsearch = Chroma.from_documents(documents, embeddings)

    chain = ConversationalRetrievalChain.from_llm(ChatOpenAI(temperature=0.3),
                                    retriever=pdfsearch.as_retriever(search_kwargs={"k": 1}),
                                    return_source_documents=True)
    return chain

# Function to generate a response based on the chat history and query
def generate_response(history, query, pdf_upload):
    global COUNT, N, chat_history, chain
    if not pdf_upload:
        raise gr.Error(message='Upload a PDF')

    if COUNT == 0:
        chain = process_file(pdf_upload)
        COUNT += 1

    # Replace with your LangChain logic to generate a response 
    result = chain({"question": query, 'chat_history': chat_history}, return_only_outputs=True)  
    chat_history += [(query, result["answer"])]
    N = list(result['source_documents'][0])[1][1]['page']  # Adjust as needed

    for char in result['answer']:
        history[-1][-1] += char  
    return history, ''  

# Function to render a specific page of a PDF file as an image
def render_file(file):
    global N
    doc = fitz.open(file.name)
    page = doc[N]
    pix = page.get_pixmap(matrix=fitz.Matrix(300/72, 300/72)) 
    image = Image.frombytes('RGB', [pix.width, pix.height], pix.samples)
    return image

# Function to render initial content from the PDF
def render_first(pdf_file): 
    # Replace with logic to process the PDF and generate an initial image
    image = Image.new('RGB', (600, 400), color = 'white') # Placeholder
    return image

# Streamlit & Gradio Interface

st.title("PDF-Powered Chatbot") 

with st.container():          
    gr.Markdown("""         
    <style>             
    .image-container { height: 680px; }         
    </style>         
    """)      

    with gr.Blocks() as demo: # Introduce a Blocks context        
        with gr.Row():                    
            enable_box = gr.Textbox(placeholder='Enter OpenAI API key',                                     
                                    show_label=False, interactive=True)                     
            disable_box = gr.Textbox(value='OpenAI API key is Set', interactive=False)                     
            change_api_key = gr.Button('Change Key')     

        with gr.Row():                    
            chatbot = gr.Chatbot(value=[], elem_id='chatbot')                      
            show_img = gr.Image(label='Upload PDF')  

            # Create multiple PDF upload buttons
            pdf_upload1 = gr.UploadButton("📁 Upload PDF 1", file_types=[".pdf"])  
            pdf_upload2 = gr.UploadButton("📁 Upload PDF 2", file_types=[".pdf"]) 
            pdf_upload3 = gr.UploadButton("📁 Upload PDF 3", file_types=[".pdf"])

        # Event handlers (adjust how you use the uploads in these functions)        
        enable_box.submit(fn=set_apikey, inputs=[enable_box], outputs=[disable_box])                 
        change_api_key.click(fn=enable_api_box, outputs=[enable_box])                 

        # Assuming you want to process the first PDF initially
        pdf_upload1.upload(fn=render_first, inputs=[pdf_upload1], outputs=[show_img]) 

        txt = gr.Textbox(label="Enter your query", placeholder="Ask a question...")         
        submit_btn = gr.Button('Submit')

        # Event handler for submit button
        @submit_btn.capture()
        def on_submit():
            add_text(chatbot, txt)
            generate_response(chatbot, txt, pdf_upload1)
            render_file(pdf_upload1)

if __name__ == "__main__":        
    gr.Interface(                
        [render_first, add_text, generate_response, render_file],                
        [pdf_upload1, chatbot, txt, pdf_upload2, pdf_upload3],               
        [show_img, chatbot, txt], # Assuming you want to display initially                
        title="PDF-Powered Chatbot"        
    ).launch()