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

from dotenv import load_dotenv, find_dotenv
from document_model import Listing
from openai_ import handle_user_prompt

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

        ###
        prompt = """
        I want to stay in a place that's warm and friendly, 
        and not too far from resturants, can you recommend a place? 
        Include a reason as to why you've chosen your selection.
        """
        result = handle_user_prompt(prompt, db, collection)
        ###

        del os.environ["OPENAI_API_KEY"]
        
        """
        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 result

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()