import os
from langchain.llms import OpenAI
import streamlit as st
import gradio as gr




#Pdf Loader
from langchain.document_loaders import PyPDFLoader
from langchain.vectorstores import Chroma
from langchain.embeddings import OpenAIEmbeddings

from langchain.agents.agent_toolkits import (
    create_vectorstore_agent,
    VectorStoreToolkit,
    VectorStoreInfo

)

os.environ['OPENAI_API_KEY'] = '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'
llm = OpenAI(temperature=0.4)

embeddings = OpenAIEmbeddings()
loader = PyPDFLoader('2022_02_11_Montelogo  Ida.pdf') # Baadme this will be updated from the Input @samunder
pages = loader.load_and_split()
store  = Chroma.from_documents(pages,embeddings,collection_name='Montelogo_Ida')


vectorstore_info = VectorStoreInfo(
    name = "Montelogo_Ida",
    description  = "Gpt on Montelogo Ida Pdf",
    vectorstore = store

)
toolkit = VectorStoreToolkit(vectorstore_info=vectorstore_info)

agent_executor = create_vectorstore_agent(
    llm=llm,
    toolkit=toolkit,
    verbose=True
)

chat_history = []

def generate_response(prompt, chat_history):
    if prompt:
        response = agent_executor.run(prompt)
        chat_history.append((prompt, response))

        return "",chat_history


with gr.Blocks() as demo:
    gr.Markdown("Upload a document and Ask Questions")
    with gr.Tab("Ask AI for cases based on Montelogo_Ida"):
      gr.Markdown("This answers based on this ")
      chatbot = gr.Chatbot()
      msg = gr.Textbox("Chat with me ❤️")
      clear = gr.ClearButton([msg, chatbot])

      msg.submit(generate_response, [msg, chatbot], [msg, chatbot])
    with gr.Tab("Upload ANY Pdf and Ask AI"):
        gr.Markdown("Work in Progress")

        # with gr.Row():
        #     # file_input = gr.File()
        #     with gr.Column(scale=0.15, min_width=0):
        #       btn = gr.UploadButton("📁", file_types=["pdf", "image", "video"])
        #       btn.upload()
        #     pdf_path = '/content/2022_02_11_Montelogo  Ida.pdf'
        #     pdf_viewer = gr.FileViewer(pdf_path)

            # image_output = gr.Image()
        # image_button = gr.Button("Flip")

    # with gr.Accordion("Open for More!"):
    #     gr.Markdown("Look at me...")

# iface = gr.Interface(
#     fn=generate_response,
#     inputs=gr.inputs.Textbox(label="User Message", placeholder="Enter your message..."),
#     outputs=gr.outputs.Textbox(label="Generated AI Response"),
#     title="Legal Case Chat",
#     description="Enter your legal case context and get AI-generated responses.",
#     theme="compact",
#     layout="vertical",
#     width="auto",
#     height=400,
#     show_tips=False
# )

demo.launch()