import gradio as gr
import Stemmer  # from PyStemmer
import bm25s
import numpy as np
import time
import os

examples = ["Can officers always order a passenger out of a car?","Find me briefs about credential searches", "Can police search an impounded car without a warrant?", "State is arguing State v. Carty is not good law"]

def show_user_query(user_message, history):
    return "", history + [[user_message, None]]

def retrieve(history):
    query = history[-1][0]
    print ("query", query)
    query_tokens = bm25s.tokenize(query, stemmer=stemmer)

    results, scores = retriever.retrieve(query_tokens, k=10)
    query_response = " ".join([str(i) for i in results])
    #query_response = retrieve(history[-1][0])
    #print (query_response)
    history[-1][1] = ""
    for character in query_response:
        history[-1][1] += character
        time.sleep(0.0001)
        yield history



with gr.Blocks() as demo:
    #gr.LoginButton()
    chatbot = gr.Chatbot()
    query_textbox = gr.Textbox()
    examples = gr.Examples(examples, query_textbox)
    next_button = gr.ClearButton(chatbot, value="Clear", visible=True)

    query_textbox.submit(show_user_query, [query_textbox, chatbot], [query_textbox, chatbot], queue=False).then(retrieve, chatbot, chatbot)


def load_bm25():
    stemmer = Stemmer.Stemmer("english")
    retriever = bm25s.BM25.load("NJ_index_LLM_chunking", mmap=False)
    return retriever, stemmer

retriever, stemmer = load_bm25()
#demo.launch(share=True, auth=(os.environ.get("username"), os.environ.get("password")))
#demo.launch(auth=("admin", "pass1234"))

demo.launch()