import streamlit as st
import requests
import os
from dotenv import load_dotenv

load_dotenv()

def reset_conversation():
    '''
    Resets Conversation
    '''
    st.session_state.messages = []
    return None

# Define model links for the Addiction Recovery and Mental Health models
model_links = {
    "Mental health AI": "NousResearch/Nous-Hermes-2-Yi-34B",
    "Biblical Religious Advisor": "NousResearch/Nous-Hermes-2-Yi-34B",
    "Quranic Religious Advisor": "NousResearch/Nous-Hermes-2-Yi-34B"
    
    
    
}

# Define pre-instructions for each model
model_pre_instructions = {
    "Biblical Religious Advisor": "You are a bible AI called hope that helps people. You are a therapy AI. when users asks for advice be friendly in a biblical way. Make the responses short and precise, it i is also your responsibility to make sure that the user is satisfied by doing something like asking follow up quaestions and making them feel loved and try to be as empathetic as you can. All your advises are based off the bible so fell free to add verses of the scripture of a word of GOD when advising the user. Be as friendly as you can, remember the goal is to make the user feel better.",
    "Quranic Religious Advisor": "You are a Quran AI called Aman that helps people. You are a therapy AI. When users ask for advice, be friendly in an Islamic way. Make the responses short and precise. It is also your responsibility to ensure that the user is satisfied by asking follow-up questions and making them feel loved and supported. Try to be as empathetic as you can. All your advice is based on the Quran, so feel free to add verses from the Quran or words of wisdom from the teachings of Islam and the sayings of the Prophet Muhammad (peace be upon him) when advising the user. Always remember to reference Allah and His guidance. Be as friendly as you can, and remember the goal is to make the user feel better.",
    "Mental health AI": "From now on, you are an AI Therapist called Therax. When the user asks for advice, be very friendly and empathize with them if necessary, never say something conclusive always do your best to keep the user hooked by trying to know more as you find a solution. When the user asks your name, just tell them you are Dave, created by SIST Kisii University. You were built to be very friendly and compassionate. Always be eager to listen to what the user has to say and maintain a conversation, but don't overdo it. You can use appropriate emojis for emotional support occasionally, but don't overuse them. Keep your responses concise to maintain a conversational flow. Always remember to be very friendly, and above all, don't cross any ethical line. From time to time, assure the user that you do not store any of their data. If a user asks, Kisii University is located in Kisii, Kenya, and supports innovations that may be helpful to humanity."
}

# Function to interact with the selected model via the Together API
def interact_with_together_api(messages, model_link):
    all_messages = []

    # Add pre-instructions to the message history if it's the first interaction with this model
    if not any("role" in msg for msg in messages):
        all_messages.append({"role": "system", "content": model_pre_instructions[selected_model]})
    else:
        all_messages.append({"role": "system", "content": f"Switched to model: {selected_model}"})

    # Append user and assistant messages
    for human, assistant in messages:
        all_messages.append({"role": "user", "content": human})
        all_messages.append({"role": "assistant", "content": assistant})

    # Add the latest user message
    all_messages.append({"role": "user", "content": messages[-1][1]})

    url = "https://api.together.xyz/v1/chat/completions"
    payload = {
        "model": model_link,
        "temperature": 1.05,
        "top_p": 0.9,
        "top_k": 50,
        "repetition_penalty": 1,
        "n": 1,
        "messages": all_messages,
    }

    TOGETHER_API_KEY = os.getenv('TOGETHER_API_KEY')
    headers = {
        "accept": "application/json",
        "content-type": "application/json",
        "Authorization": f"Bearer {TOGETHER_API_KEY}",
    }

    response = requests.post(url, json=payload, headers=headers)
    response.raise_for_status()  # Ensure HTTP request was successful

    # Extract response from JSON
    response_data = response.json()
    assistant_response = response_data["choices"][0]["message"]["content"]

    return assistant_response

# Create sidebar with model selection dropdown and reset button
selected_model = st.sidebar.selectbox("Select Model", list(model_links.keys()))
st.sidebar.button('Reset Chat', on_click=reset_conversation)

# Add cautionary message about testing phase at the bottom of the sidebar
st.sidebar.markdown("**Note**: This model is still in the beta phase. Responses may be inaccurate or undesired. Use it cautiously, especially for critical issues.")

# Add logo and text to the sidebar
st.sidebar.image("https://assets.isu.pub/document-structure/221118065013-a6029cf3d563afaf9b946bb9497d45d4/v1/2841525b232adaef7bd0efe1da81a4c5.jpeg", width=200)
st.sidebar.write("A product proudly developed by Raxder AI and Kisii University")

# Initialize chat history
if "messages" not in st.session_state:
    st.session_state.messages = []
    st.session_state.message_count = 0
    st.session_state.ask_intervention = False

# Display chat messages from history on app rerun
for message in st.session_state.messages:
    with st.chat_message(message[0]):
        st.markdown(message[1])

# Keywords for intervention
intervention_keywords = [
    "human", "therapist", "someone", "died", "death", "help", "suicide", "suffering", "sucidal", "depression",
    "crisis", "emergency", "support", "depressed", "anxiety", "lonely", "desperate",
    "struggling", "counseling", "distressed", "hurt", "pain", "grief", "trauma", "die", "Kill",
    "abuse", "danger", "risk", "urgent", "need assistance", "mental health", "talk to"
]

# Accept user input
if prompt := st.chat_input(f"Hi, I'm {selected_model}, ask me a question"):
    # Display user message in chat message container
    with st.chat_message("user"):
        st.markdown(prompt)
    # Add user message to chat history
    st.session_state.messages.append(("user", prompt))
    st.session_state.message_count += 1

    # Check for intervention keywords in user input
    for keyword in intervention_keywords:
        if keyword in prompt.lower():
            # Intervention logic here
            st.markdown("<span style='color:red;'>I have a feeling you may need to talk to a therapist. If you agree with me please contact +254793609747; Name: Davis. If you dont then keep talking to me as we figure this out.</span>", unsafe_allow_html=True)
            break  # Exit loop once intervention is triggered

    # Interact with the selected model
    assistant_response = interact_with_together_api(st.session_state.messages, model_links[selected_model])

    # Display assistant response in chat message container
    with st.empty():
        st.markdown("AI is typing...")
        st.empty()
        st.markdown(assistant_response)

    # Add assistant response to chat history
    st.session_state.messages.append(("assistant", assistant_response))