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

#from dotenv import load_dotenv, find_dotenv
from custom_utils import connect_to_database, rag_ingestion, handle_user_prompt

#from pydantic import BaseModel
#from typing import Optional

#from IPython.display import display, HTML

lock = threading.Lock()

#_ = load_dotenv(find_dotenv())

RAG_INGESTION = True

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:
        db, collection = connect_to_database()

        if (RAG_INGESTION):
            rag_ingestion(collection)

        """
        print("777")
        search_path = "address.country"

        print("888")
        # Create a match stage
        match_stage = {
            "$match": {
               search_path: re.compile(r"United States"),
               "accommodates": { "$gt": 1, "$lt": 3}
            }
        }

        print("999")
        additional_stages = [match_stage]
        """

        print("000")
        #result = handle_user_query(openai_api_key, query, db, collection, additional_stages)
        return handle_user_prompt(openai_api_key, prompt, db, collection)

gr.close_all()

PROMPT = "I want to stay in a place that's modern and clean, walking distance from restaurants."

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

demo.launch()