# Importing required packages
import streamlit as st
import uuid
from huggingface_hub import InferenceClient

# Initialize the HuggingFace inference client
def initialize_client(api_key):
    return InferenceClient(model="HuggingFaceH4/zephyr-7b-beta", token=api_key)

INIT_PROMPT = """
\n\nHuman: You are DarijaBot, a helpful assistant that can converse in both Latin and Arabic alphabet Darija. You will help the users learn about Wardley Mapping.
"""

TRAINING_PROMPT = """
Here is an outline for a training course that you will give to the user. It covers the key principles of Wardley Mapping:
"""

INTRO_PROMPT = """
Hello, I'm DarijaBot! I can help you with your questions about Wardley Mapping in Darija. You can type in either Latin or Arabic script.
"""

REG_PROMPT = """
\n\nHuman: Here is the user's question in Darija:
<question>
{QUESTION}
</question>
\n\nAssistant: [DarijaBot] <response>
"""

new_prompt = []

if "session_id" not in st.session_state:
    st.session_state.session_id = str(uuid.uuid4())

st.set_page_config(page_title="DarijaBot - ChatBot")
st.sidebar.title("DarijaBot - ChatBot")
st.sidebar.title("Wardley Mapping Mentor")
st.sidebar.divider()
st.sidebar.markdown("Developed by [Bahae Eddine HALIM](https://linkedin.com/in/halimbahae)", unsafe_allow_html=True)
st.sidebar.markdown("Current Version: 0.0.4")
st.sidebar.markdown("Using HuggingFaceH4/zephyr-7b-beta API")
st.sidebar.markdown(st.session_state.session_id)
st.sidebar.divider()

# Prompt the user for the API key
user_huggingface_api_key = st.sidebar.text_input("Enter your HuggingFace API Key:", placeholder="hf_...", type="password")

if "messages" not in st.session_state:
    st.session_state["messages"] = []
    st.session_state.messages.append({"role": "assistant", "content": INTRO_PROMPT})

if "all_prompts" not in st.session_state:
    st.session_state["all_prompts"] = INIT_PROMPT + TRAINING_PROMPT

if user_huggingface_api_key:
    client = initialize_client(user_huggingface_api_key)
else:
    st.warning("Please enter your HuggingFace API key", icon="⚠️")

for message in st.session_state.messages:
    if message["role"] in ["user", "assistant"]:
        with st.chat_message(message["role"]):
            new_prompt.append(message["content"])
            st.markdown(message["content"])

if user_huggingface_api_key:
    if user_input := st.chat_input("How can I help with Wardley Mapping? (Darija: Latin or Arabic script)"):
        prompt = REG_PROMPT.format(QUESTION=user_input)
        st.session_state.all_prompts += prompt
        st.session_state.messages.append({"role": "user", "content": user_input})
        with st.chat_message("user"):
            st.markdown(user_input)
        with st.chat_message("assistant"):
            message_placeholder = st.empty()
        full_response = ""
        try:
            response = client.chat_completion(
                [{"role": "system", "content": "You are DarijaBot, a helpful assistant that can converse in both Latin and Arabic alphabet Darija."}, {"role": "user", "content": user_input}],
                max_tokens=512,
                temperature=0.7,
                top_p=0.95
            ).choices[0].message.content
            full_response += response
            message_placeholder.markdown(full_response + "▌")
            message_placeholder.markdown(full_response)
        except Exception as e:
            st.error(f"Error: Unable to fetch response from the API.\nDetails: {e}")
        st.session_state.messages.append({"role": "assistant", "content": full_response})
        st.session_state.all_prompts += full_response