# Setup

# Import the necessary Libraries
import os
import json
import uuid
import gradio as gr
#import tiktoken
from datasets import load_dataset

#import pandas as pd

from openai import OpenAI

from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_core.documents import Document

from langchain_community.embeddings.sentence_transformer import (
    SentenceTransformerEmbeddings
)
from langchain_community.vectorstores import Chroma

#from google.colab import userdata, drive

from langchain_community.document_loaders import PyPDFDirectoryLoader

#from google.colab import userdata



from huggingface_hub import CommitScheduler
from pathlib import Path




# Create Client
client = OpenAI(
    base_url="https://api.openai.com/v1",
    api_key=os.environ['CarlosGM']
)


#api_key = os.environ.get("CarlosGM")

#client = OpenAI(api_key=api_key)





model_name = 'gpt-3.5-turbo'


# Define the embedding model and the vectorstore

embedding_model = SentenceTransformerEmbeddings(model_name='thenlper/gte-large')

# Load the persisted vectorDB

## persisted_vectordb_location = '/content/drive/MyDrive/dataset_db'

dataset_10k_collection = 'Dataset-IBM-Meta-aws-google-msft'

vectorstore_persisted = Chroma(
    collection_name=dataset_10k_collection,
    persist_directory= './dataset_db',
    embedding_function=embedding_model
)


retriever = vectorstore_persisted.as_retriever(
    search_type='similarity',
    search_kwargs={'k': 5}
)


# Prepare the logging functionality

log_file = Path("logs/") / f"data_{uuid.uuid4()}.json"
log_folder = log_file.parent

scheduler = CommitScheduler(
    repo_id="10k-logs",
    repo_type="dataset",
    folder_path=log_folder,
    path_in_repo="data",
    every=2
)

# Define the Q&A system message

qna_system_message = """
You are an assistant to a financial services firm who answers user queries on annual 10 K reports.
User input will have the context required by you to answer user questions.
This context will begin with the token: ###Context.
The context contains references to specific portions of a document relevant to the user query.
The source for a context will begin with the token ###Source

User questions will begin with the token: ###Question.



Please answer only using the context provided in the input. Do not mention anything about the context in your final answer.

Please adhere to the following guidelines:
- Your response should only be about the question asked and nothing else.
- Answer only using the context provided.
- Do not mention anything about the context in your final answer.
- If the answer is not found in the context, it is very very important for you to respond with "I don't know. Please check the docs @ 'Dataset-10k file'"
- Always quote the source when you use the context. Cite the relevant source at the end of your response under the section - Source:
- Do not make up sources. Use the files provided in the sources section of the context and nothing else. You are prohibited from providing other sources.


If the answer is not found in the context, respond "I don't know".

Here is an example of how to structure your response:

Answer:
[Answer]

Source:
[Source]

"""




# Define the user message template

qna_user_message_template = """
###Context
Here are some documents that are relevant to the question.
{context}

###Question
{question}
"""


# Define the predict function that runs when 'Submit' is clicked or when a API request is made

def predict(user_input,company):

    #filter = "dataset/"+company+"-10-k-2023.pdf"
    #relevant_document_chunks = vectorstore_persisted.similarity_search(user_input, k=5, filter={"source":filter})

    # Create context_for_query

    relevant_document_chunks = retriever.invoke(user_input)
    
    context_list = [d.page_content for d in relevant_document_chunks]
    context_for_query = ". ".join(context_list)
   

    # Create messages

    prompt = [
        {'role':'system', 'content': qna_system_message},
        {'role': 'user', 'content': qna_user_message_template.format(        
            context=context_for_query,
            question=user_input
            )
        }
    ]


    # Get response from the LLM

    try:
        response = client.chat.completions.create(
            model=model_name,
            messages=prompt,
            temperature=0
        )

        prediction = response.choices[0].message.content.strip()
    except Exception as e:
        prediction = f'Sorry, I encountered the following error: \n {e}'








    # While the prediction is made, log both the inputs and outputs to a local log file

    # While writing to the log file, ensure that the commit scheduler is locked to avoid parallel
    # access

    with scheduler.lock:
        with log_file.open("a") as f:
            f.write(json.dumps(
                {
                    'user_input': [user_input, company],
                    'retrieved_context': context_for_query,
                    'model_response': prediction
                }
            ))
            f.write("\n")

    return prediction

# Set-up the Gradio UI



# Add text box and radio button to the interface
# The radio button is used to select the company 10k report in which the context needs to be retrieved.

textbox = gr.Textbox(placeholder='Enter your query here', lines=6)
company = gr.Radio(['aws', 'google', 'ibm', 'meta', 'msft'], label= "Select Company 10-k Report")

# Create the interface

demo = gr.Interface(
    fn=predict,
    inputs=[textbox,company],
    outputs= "text",
    title= "10-k Report Q&A",
    description = "This Web API presents an inteface to ask questions about the 10-k reports"
)
    

# For the inputs parameter of Interface provide [textbox,company]


demo.queue()
demo.launch()