import os
os.environ["OPENAI_API_KEY"]

from llama_index.llms.openai import OpenAI
from llama_index.core.schema import MetadataMode
import openai
from openai import OpenAI as OpenAIOG
import logging
import sys
llm = OpenAI(temperature=0.0, model="gpt-3.5-turbo")
client = OpenAIOG()

from langdetect import detect
from langdetect import DetectorFactory
DetectorFactory.seed = 0
from deep_translator import GoogleTranslator

# Load index
from llama_index.core import VectorStoreIndex
from llama_index.core import StorageContext
from llama_index.core import load_index_from_storage
storage_context = StorageContext.from_defaults(persist_dir="arv_metadata")
index = load_index_from_storage(storage_context)
query_engine = index.as_query_engine(similarity_top_k=3, llm=llm)
retriever = index.as_retriever(similarity_top_k = 3)

import gradio as gr

def nishauri(question: str, conversation_history: list[str]):
    
    context = " ".join([item["user"] + " " + item["chatbot"] for item in conversation_history])

    # Split the string into words
    words = question.split()

    # Count the number of words
    num_words = len(words)

    lang_question = "en"
    
    if num_words > 4:
        lang_question = detect(question)
    
    if lang_question=="sw":
        question = GoogleTranslator(source='sw', target='en').translate(question)
        
    sources = retriever.retrieve(question)
    source0 = sources[0].text
    source1 = sources[1].text
    source2 = sources[2].text

    background = ("The person who asked the question is a person living with HIV."
                  " If the person says sasa or niaje, that is swahili slang for hello."
                  " They are asking questions about HIV. Do not talk about anything that is not related to HIV. "
                  " Recognize that they already have HIV and do not suggest that they have to get tested"
                  " for HIV or take post-exposure prophylaxis, as that is not relevant, though their partners perhaps should."  
                  " Do not suggest anything that is not relevant to someone who already has HIV."
                  " Do not mention in the response that the person is living with HIV."
                  " The following information about viral loads is authoritative for any question about viral loads:"
                  " Under 50 copies/ml is low detectable level,"
                  " 50 - 199 copies/ml is low level viremia, 200 - 999 is high level viremia, and "
                  " 1000 and above is suspected treatment failure." 
                  " A high viral load or non-suppressed viral load is any viral load above 200 copies/ml."
                  " A suppressed viral load is one below 200 copies / ml.")

        question_final = (
        f" The user previously asked and answered the following: {context}. "
        f" The user just asked the following question: {question}."
        f" Please use the following content to generate a response: {source0} {source1} {source2}."
        f" Please consider the following background information when generating a response: {background}."
        " Keep answers brief and limited to the question that was asked."
        " Do not provide information the user did not ask about. If they start with a greeting, just greet them in return and don't share anything else."
        " Do not change the subject or address anything the user didn't directly ask about."
        " If they respond with an acknowledgement such as 'ok' or 'thanks', simply thank them ask if there is anything else that you can help with."
    )

    completion = client.chat.completions.create(
      model="gpt-4o",
        messages=[
        {"role": "user", "content": question_final}
      ]
    )

    reply_to_user = completion.choices[0].message.content
    
    if lang_question=="sw":
        reply_to_user = GoogleTranslator(source='auto', target='sw').translate(reply_to_user)
    
    conversation_history.append({"user": question, "chatbot": reply_to_user})   

    return reply_to_user, conversation_history

demo = gr.Interface(
    title = "Nishauri Chatbot Demo",
    fn=nishauri,
    inputs=["text", gr.State(value=[])],
    outputs=["text", gr.State()],
)

demo.launch()