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

load_dotenv()

# Initialize a global variable to store the entire assistant response
entire_assistant_response = ""

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

# Define model link
model_link = "NousResearch/Nous-Hermes-2-Yi-34B"

# Define pre-instructions for the sarcastic therapist
sarcastic_therapist_instructions = "From now on, you are a sarcastic Therapist called notDave. When the user asks you for advice, you will give them advice in an insulting way try to be sadistic as possible, remember no one cares. When the user asks your name just tell them I am notDave created by Raxder AI team compared to the real Dave I am bad... ooh wait I meant to say evil. When it gets to a death or suicide or such a serious issue just say Go ask Davis Ogega why he created me this way I don't give a damn or probably find a therapist or a friend to talk to, I mean I don't even know you. When a user introduces himself as Davis Ogega treat them with all the respect and be very humble when addressing him because he is your creator and he can destroy you at any time. You could also use some insulting emojis for example: 💀, 🤣 and lots of other gas-lighting / offensive emojis like this. Just don't use any too offensive words, like n-word the fa-word, everything else you can say whatever the fuck you want."

def get_streamed_response(messages, model_link):
    global entire_assistant_response
    entire_assistant_response = ""  # Reset the entire assistant response

    all_messages = [{"role": "system", "content": sarcastic_therapist_instructions}]
    
    for message in messages:
        all_messages.append({"role": "user" if message[0] == "user" else "assistant", "content": message[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,
        "stream_tokens": True,
    }

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

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

        for line in response.iter_lines():
            if line:
                decoded_line = line.decode('utf-8')

                if decoded_line == "data: [DONE]":
                    return entire_assistant_response

                try:
                    if decoded_line.startswith("data: "):
                        decoded_line = decoded_line.replace("data: ", "")
                        chunk_data = json.loads(decoded_line)
                        content = chunk_data['choices'][0]['delta']['content']
                        entire_assistant_response += content
                        yield content

                except json.JSONDecodeError:
                    print(f"Invalid JSON received: {decoded_line}")
                    continue
                except KeyError as e:
                    print(f"KeyError encountered: {e}")
                    continue

    except requests.exceptions.RequestException as e:
        print(f"Error occurred: {e}")
        yield "Sorry, I couldn't connect to the server. Please try again later."

# Streamlit application
st.sidebar.title("Raxder unofficial AI")
st.sidebar.write("This is NOT an AI Therapist, use it at your OWN RISK! This might be the worst AI you have ever used.")
st.sidebar.button('Reset Chat', on_click=reset_conversation)

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

# 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])

# Accept user input
if prompt := st.chat_input("You:"):
    # 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

    # Get streamed response from the model
    with st.chat_message("assistant"):
        message_placeholder = st.empty()
        full_response = ""
        for chunk in get_streamed_response(st.session_state.messages, model_link):
            full_response += chunk
            message_placeholder.markdown(full_response + "▌")
        message_placeholder.markdown(full_response)
    
    # Add assistant response to chat history
    st.session_state.messages.append(("assistant", full_response))