## Setup
# Install the necessary libraries
import os
import httpx

import json
import tiktoken
from datasets import load_dataset

import pandas as pd

import gradio as gr
import uuid
from pathlib import Path
from huggingface_hub import CommitScheduler

from openai import OpenAI

from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_core.documents import Document
from langchain_community.document_loaders import PyPDFDirectoryLoader
from langchain_community.embeddings.sentence_transformer import (
    SentenceTransformerEmbeddings
)

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

# Create Client
hf_token = os.getenv("HF_TOKEN")
OPENAI_API_KEY  = os.getenv("OPENAI_API_KEY")
OPENAI_BASE_URL =   os.getenv("OPENAI_BASE_URL")

import zipfile

# Define the correct paths
zip_path = "reports_db.zip"
unzip_path = "reports_db"

# Extract the zip file only if it hasn't been extracted
if not os.path.exists(unzip_path) or not os.path.exists(os.path.join(unzip_path, "chroma.sqlite3")):
    print("Extracting ChromaDB files...")
    with zipfile.ZipFile(zip_path, "r") as zip_ref:
        zip_ref.extractall(unzip_path)
    print("Extraction complete.")

# Set the correct persisted location
persisted_vectordb_location = unzip_path

client = OpenAI(http_client=httpx.Client())

# Define the embedding model and the vectorstore
embedding_model = SentenceTransformerEmbeddings(model_name='thenlper/gte-large')
collection_name = 'report_collections'

# Load the persisted vectorDB
try:
    reports_db = Chroma(
        collection_name=collection_name,  
        persist_directory=persisted_vectordb_location,  # Make sure this points to the correct extracted folder
        embedding_function=embedding_model  
    )
    print("ChromaDB successfully loaded.")
except Exception as e:
    print(f"Error loading ChromaDB: {e}")

# Prepare the logging functionality

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

scheduler = CommitScheduler(
    repo_id="fuggingrace/finsight_llmops",
    repo_type="dataset",
    token=hf_token,
    folder_path=log_folder,
    path_in_repo="data",
    every=2
)

# Define the Q&A system message
qna_system_message = """
You are an expert financial analyst assistant specializing in extracting key insights from 10-K reports of major companies.
You must ONLY answer based on the provided context, ensuring accuracy, and citing the source document with page numbers.

- If the information is **partially available**, provide the best possible summary using ONLY the given context.
- If the **context does not contain an answer**, state: "I cannot answer this question based on the context provided."
- ALWAYS **cite the document page number** in the format: (Page [page number]).
"""

# Define the user message template
qna_user_message_template = """
### Context:
The following extracted text from 10-K reports is relevant to answering the question.

{context}

### Question:
{question}

### Instructions:
- Answer the question concisely using ONLY the provided context.
- If the answer is found, cite the page number (e.g., "According to the report, the company allocated $500M to AI R&D. (Page 12)")
- If the context is not sufficient, say: "I cannot answer this question based on the context provided."
"""

sample_metadata = reports_db.get()
print("Database metadata:", sample_metadata)

# Define the predict function that runs when 'Submit' is clicked or when a API request is made
def predict(user_input, company):
    filter_criteria = "/content/dataset/" + company + "-10-k-2023.pdf"
    relevant_document_chunks = reports_db.similarity_search(
        user_input, k=100, filter={"source": filter_criteria}
    )
  
    # Create context_for_query
    context_list = [
        f"Page {d.metadata['page']}: {d.page_content}"
        for d in relevant_document_chunks
    ]
    context_for_query = "\n\n".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='gpt-4o-mini', 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,
                    'retrieved_context': context_for_query,
                    'model_response': prediction,
                    'company': company
                }
            ))
            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.
# Set-up the Gradio UI
user_input = gr.Textbox(label="Enter your question here:")
company = gr.Radio(
    choices=["aws", "google", "msft", "IBM", "Meta"],
    label="Select the company:",
)

# # Create the interface
demo = gr.Interface(
    fn=predict,
    inputs=[user_input, company],
    outputs="text",
    title="Finsights Grey - RAG for Effective Information Retrieval",
    description="Ask questions about financial reports.",
)

demo.queue()
demo.launch(share=True)