# Import librairies 
from pathlib import Path
import sys
import os 
import openai
import llama_index
from llama_index import SimpleDirectoryReader, GPTListIndex, readers, LLMPredictor, PromptHelper, ServiceContext, GPTVectorStoreIndex, StorageContext, load_index_from_storage, download_loader,  GPTRAKEKeywordTableIndex
from llama_index.retrievers import VectorIndexRetriever
from langchain import OpenAI
from llama_index.node_parser import SimpleNodeParser
import gradio as gr
from llama_index.optimization.optimizer import SentenceEmbeddingOptimizer
from langchain.chat_models import ChatOpenAI
from llama_index.readers import Document
import io
from PyPDF2 import PdfReader
from azure.storage.filedatalake import DataLakeServiceClient
from llama_index.indices.vector_store.base import GPTVectorStoreIndex
from adlfs import AzureBlobFileSystem
import time 

# Blob storage parameters
account_name = 'apeazdlkini07s'
account_key = os.environ['account_key']
file_system_name = "gpt"
service_client = DataLakeServiceClient(account_url=f"https://{account_name}.dfs.core.windows.net", credential=account_key)
file_system_client = service_client.get_file_system_client(file_system_name)

AZURE_ACCOUNT_NAME = account_name
AZURE_ACCOUNT_KEY = account_key

assert AZURE_ACCOUNT_NAME is not None and AZURE_ACCOUNT_NAME != ""

fs = AzureBlobFileSystem(account_name=AZURE_ACCOUNT_NAME, account_key=AZURE_ACCOUNT_KEY)


def construct_index(doc):

  ## Define the prompt helper
    # Set maximum input size
  max_input_size = 400

    # Set number of output tokens
  num_output = 400 # About 300 words

    #Set the chunk size limit 
  chunk_size_limit = 600 # About 450 words ~ 1 page

    # Set maximum chunk overlap
  max_chunk_overlap = 1

    # Set chunk overlap ratio 
  chunk_overlap_ratio = 0.5

    # Define prompt helper
  prompt_helper = PromptHelper(max_input_size, num_output, max_chunk_overlap, chunk_size_limit, chunk_overlap_ratio)

  ## Define the LLM predictor 
  llm_predictor = LLMPredictor(llm=ChatOpenAI(temperature=0.4, model_name="gpt-4-32k", max_tokens=num_output))
  
  ## Define Service Context
  service_context = ServiceContext.from_defaults(llm_predictor=llm_predictor, prompt_helper=prompt_helper)

  ## Indexation process and saving in the disk
  index = GPTVectorStoreIndex.from_documents(doc, service_context=service_context)

  # save index to disk
  index.set_index_id("vector_index")

  return index

def extract_text(file):
    # Open the PDF file in binary mode
    with open(file.name, 'rb') as f:
        # Initialize a PDF file reader object
        pdf_reader = PdfReader(f)

        # Initialize an empty string for storing the extracted text
        text = ''

        # Loop through the number of pages
        for page in pdf_reader.pages:
            # Add the text from each page to the text string
            text += page.extract_text()

        return text, os.path.basename(file.name)
  
def ask_ai(doc, question):
  
  text, file_name = extract_text(doc)
  index = construct_index([Document(text)])
  
  # Save index to Azure blob storage
  index.storage_context.persist(f'gpt/storage_demo/{file_name}', fs=fs)

  # Rebuild storage context
  storage_context = StorageContext.from_defaults(persist_dir=f'gpt/storage_demo/{file_name}', fs=fs)

  # Load index
  index = load_index_from_storage(storage_context)
  
  # Define the query & the querying method
  query_engine = index.as_query_engine(optimizer=SentenceEmbeddingOptimizer(percentile_cutoff=0.8)) 
  query = 'Your task is to  answer a question on the report loaded and give insights to an investment team in Infrastructure. Make your response as clear and precise as possible. The question is:' + str(question)
  response = query_engine.query(query)
  
  # Display the chunks retrieved to produce the response
  sources = []
  for node in response.source_nodes:
    node_text_start= 'START: ' + node.node.text.strip().replace('\n', ' ')[:100]
    node_text_end = 'END: ' + node.node.text.strip().replace('\n', ' ')[-100:]
    sources.append((node_text_start, node_text_end))

  return response.response

header = """<center><b><p style=\"color: #E13C32; font-size: 36px;\">My Ardian Chatbot</p></b></center>
<i><p style=\"font-size: 16px; color: grey;\">Please make sure to formulate clear and precise questions and to add contextual information when possible. This will help the tool produce the most relevant response. Adopt an iterative approach and ask for more details or explanations when necessary.</br><i/></p>"""

footnote = "<p style=\"font-size: 16px; color: grey;\"> ⚠ The chatbot doesn't have a memory, it doesn't remember what it previously generated.</a></p>"

theme = gr.themes.Base(
  primary_hue="red",
  secondary_hue="gray",
  font=['FuturaTOT', '=', '36px']
)

with gr.Blocks(theme=theme) as demo:
    gr.Markdown(header)

    download_button = gr.inputs.File(label="Upload a PDF")
    chatbot = gr.Chatbot()
    question = gr.Textbox(label='Question', info="Please write your question here.")
    clear = gr.Button("Clear")
    
    def respond(message, chat_history, doc):
        bot_message = ask_ai(doc, message)
        chat_history.append((message, bot_message))
        time.sleep(2)
        return "", chat_history

    question.submit(respond, [question, chatbot, download_button], [question, chatbot])
    
    clear.click(lambda: None, None, chatbot, queue=False)
    gr.Markdown(footnote)

demo.launch(auth=(os.environ['username'],os.environ['password']))