import os
import gradio as gr
from gradio_client import Client

MY_HF_TOKEN_KEY = os.environ["MY_HF_TOKEN_KEY"]

#iface = gr.load(name="Ghana-NLP/Khaya-chat-bot",hf_token=MY_HF_TOKEN_KEY,src ='spaces')
client = Client("Ghana-NLP/Khaya-chat-backend",hf_token=MY_HF_TOKEN_KEY)

def generate(prompt,language,history):
    history.append("USER: "+prompt+"\n-----\n")
    output = client.predict(prompt,language)
    history.append("KHAYA: "+output+"\n*****\n")
    full_output = ""
    for el in history:
        full_output = full_output + el
    return  history,full_output

with gr.Blocks() as demo:
    history_var = gr.State([]) # conversation history
    title = gr.Markdown(
    """
    # African Language Chatbot (Khaya)
    1. SELECT YOUR LANGUAGE
    2. TYPE IN SELECTED LANGUAGE!!
    3. REFRESH PAGE TO START FRESH CONVERSATION \n\n
    Based on the <a href="https://translation.ghananlp.org/">Khaya AI Translation API</a>, <a href="https://lesan.ai/">Lesan AI</a>, Google Translate API and the Mixtral model. Lesan is used for <b>Amharic</b> and <b>Tigrinya</b>, <a href="https://translation.ghananlp.org/">Khaya AI</a> is used for <b>Twi</b>, <b>Dagbani</b>, <b>Ewe</b>, <b>Ga</b>, <b>Gurene</b>, <b>Fante</b>, <b>Kikuyu</b>, <b>Kimeru</b>, <b>Luo</b>, <b>Yoruba</b> and Google is used for <b>Hausa</b>, <b>Swahili</b> and <b>Shona</b>.
    """)    
    language_selector = gr.Dropdown(["Amharic","Twi","Dagbani","Ewe","Ga","Gurene","Fante","Hausa", "Kikuyu", "Kimeru", "Luo","Shona","Swahili","Tigrinya","Yoruba"],value="Twi",label="Choose Language! (default is Twi)", info="Language:")
    with gr.Row():
        output = gr.Text(label="OUTPUT WINDOW (Conversation so far):")
    with gr.Row():
        prompt = gr.Text(label="Enter Text In Your Local Language:")
    btn = gr.Button("Chat")
    btn.click(generate, inputs=[prompt,language_selector,history_var], outputs=[history_var,output])

if __name__ == "__main__":
    demo.queue(max_size=20).launch()