import gradio as gr
import logging, os, sys, threading
import pandas as pd

from dotenv import load_dotenv, find_dotenv
from datasets import load_dataset
from document_model import Listing

lock = threading.Lock()

_ = load_dotenv(find_dotenv())

RAG_INGESTION = False # load, split, embed, and store documents

RAG_OFF        = "Off"
RAG_NAIVE  = "Naive RAG"
RAG_ADVANCED = "Advanced RAG"

logging.basicConfig(stream = sys.stdout, level = logging.INFO)
logging.getLogger().addHandler(logging.StreamHandler(stream = sys.stdout))

def invoke(openai_api_key, prompt, rag_option):
    if not openai_api_key:
        raise gr.Error("OpenAI API Key is required.")
    if not prompt:
        raise gr.Error("Prompt is required.")
    if not rag_option:
        raise gr.Error("Retrieval-Augmented Generation is required.")

    with lock:
        os.environ["OPENAI_API_KEY"] = openai_api_key

        ###
        dataset = load_dataset("MongoDB/airbnb_embeddings", streaming=True, split="train")
        dataset = dataset.take(100)
        dataset_df = pd.DataFrame(dataset)
        dataset_df.head(5)
        print("Columns:", dataset_df.columns)
        records = dataset_df.to_dict(orient='records')
        # To handle catch `NaT` values
        for record in records:
            for key, value in record.items():
                # Check if the value is list-like; if so, process each element.
                if isinstance(value, list):
                    processed_list = [None if pd.isnull(v) else v for v in value]
                    record[key] = processed_list
                # For scalar values, continue as before.
                else:
                    if pd.isnull(value):
                        record[key] = None
        try:
          # Convert each dictionary to a Movie instance
          listings = [Listing(**record).dict() for record in records]
          # Get an overview of a single datapoint
          print(listings[0].keys())
        except ValidationError as e:
          print(e)
        ###
        
        """
        if (RAG_INGESTION):
            if (rag_option == RAG_LANGCHAIN):
                #rag = LangChainRAG()
                #rag.ingestion(config)
            elif (rag_option == RAG_LLAMAINDEX):
                #rag = LlamaIndexRAG()
                #rag.ingestion(config)
    
        try:
                #rag = LangChainRAG()
                #completion, callback = rag.rag_chain(config, prompt)
                #result = completion["result"]
            elif (rag_option == RAG_LLAMAINDEX):
                #rag = LlamaIndexRAG()
                #result, callback = rag.retrieval(config, prompt)
            else:
                #rag = LangChainRAG()
                #completion, callback = rag.llm_chain(config, prompt)
                #result = completion.generations[0][0].text
        except Exception as e:
            err_msg = e
    
            raise gr.Error(e)
        finally:
            del os.environ["OPENAI_API_KEY"]
        """

        return "TODO"

gr.close_all()

demo = gr.Interface(
    fn = invoke, 
    inputs = [gr.Textbox(label = "OpenAI API Key", type = "password", lines = 1), 
              gr.Textbox(label = "Prompt", value = "TODO", lines = 1),
              gr.Radio([RAG_OFF, RAG_NAIVE, RAG_ADVANCED], label = "Retrieval-Augmented Generation", value = RAG_ADVANCED)],
    outputs = [gr.Textbox(label = "Completion")],
    title = "Context-Aware Reasoning Application",
    description = os.environ["DESCRIPTION"]
)

demo.launch()