import numpy as np
import streamlit as st
from openai import OpenAI
import os

# inicializar o cliente
client = OpenAI(
    base_url="https://api-inference.huggingface.co/v1",
    api_key=os.environ.get('API_KEY')  # Substitua pela sua chave API
)

# Definir mensagem do sistema para configuração inicial
smessage = """
Seu nome é Assistente. 
Responda sempre em português do Brasil. 
As respostas devem ser concisas. 
O tom da conversa deve ser informal.
"""

# Criar modelos suportados e seus limites de tokens
model_links = {
    "Meta-Llama-3-8B-Instruct": ("meta-llama/Meta-Llama-3-8B-Instruct", 4096),
    "Mixtral-8x7B-Instruct-v0.1": ("mistralai/Mixtral-8x7B-Instruct-v0.1", 32768),
    "Nous-Hermes-2-Mixtral-8x7B-DPO": ("NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO", 32768),
    "Yi-1.5-34B-Chat": ("01-ai/Yi-1.5-34B-Chat", 4096),
    "Mistral-7B-Instruct-v0.1": ("mistralai/Mistral-7B-Instruct-v0.1", 8192),
    "Mistral-7B-Instruct-v0.2": ("mistralai/Mistral-7B-Instruct-v0.2", 8192),
    "Mistral-7B-Instruct-v0.3": ("mistralai/Mistral-7B-Instruct-v0.3", 8192),
    "Zephyr-7B-Beta": ("HuggingFaceH4/zephyr-7b-beta", 8192),
    "Zephyr-7B-Alpha": ("HuggingFaceH4/zephyr-7b-alpha", 8192),
    "Phi-3-mini-4k-instruct": ("microsoft/Phi-3-mini-4k-instruct", 4096),
}

def reset_conversation():
    '''
    Reinicia a Conversa
    '''
    st.session_state.conversation = []
    st.session_state.messages = []
    st.session_state.system_message_added = False
    add_system_message()
    return None

def add_system_message():
    system_message = {
        "role": "system",
        "content": smessage
    }
    if "system_message_added" not in st.session_state or not st.session_state.system_message_added:
        st.session_state.messages.append(system_message)
        st.session_state.system_message_added = True

def calculate_max_tokens(model_name, current_context_length):
    repo_id, model_limit = model_links[model_name]
    available_tokens = model_limit - current_context_length
    return max(100, min(available_tokens - 100, 3000))  # Garante um mínimo de 100 e máximo de 3000

# Definir os modelos disponíveis
models = list(model_links.keys())

# Criar a barra lateral com o menu suspenso para seleção de modelo
selected_model = st.sidebar.selectbox("Selecione o Modelo", models)

# Criar um slider de temperatura
temp_values = st.sidebar.slider('Selecione um valor de temperatura', 0.0, 1.0, (0.5))

# Adicionar botão de reiniciar para limpar a conversa
st.sidebar.button('Reiniciar Chat', on_click=reset_conversation)  # Botão de reiniciar

# Criar descrição do modelo
st.sidebar.write(f"Você está conversando com **{selected_model}**")
st.sidebar.markdown("*O conteúdo gerado pode ser impreciso ou falso.*")

if "prev_option" not in st.session_state:
    st.session_state.prev_option = selected_model

if st.session_state.prev_option != selected_model:
    st.session_state.messages = []
    st.session_state.prev_option = selected_model
    reset_conversation()

# Obter o modelo que queremos usar
repo_id = model_links[selected_model][0]

st.subheader(f'Chat com {selected_model}')

# Definir um modelo padrão
if selected_model not in st.session_state:
    st.session_state[selected_model] = repo_id

# Inicializar histórico do chat
if "messages" not in st.session_state:
    st.session_state.messages = []

# Adicionar mensagem do sistema para configuração inicial
add_system_message()

# Exibir mensagens do chat do histórico ao reiniciar o app
for message in st.session_state.messages:
    if message["role"] != "system":
        with st.chat_message(message["role"]):
            st.markdown(message["content"])

# Aceitar entrada do usuário
if prompt := st.chat_input(f"Olá, sou o {selected_model}. Faça uma pergunta"):
    # Exibir mensagem do usuário no container de mensagens do chat
    with st.chat_message("user"):
        st.markdown(prompt)
    # Adicionar mensagem do usuário ao histórico do chat
    st.session_state.messages.append({"role": "user", "content": prompt})

    # Exibir resposta do assistente no container de mensagens do chat
    with st.chat_message("assistant"):
        try:
            current_context_length = sum(len(m['content']) for m in st.session_state.messages)
            max_tokens = calculate_max_tokens(selected_model, current_context_length)

            stream = client.chat.completions.create(
                model=repo_id,
                messages=[
                    {"role": m["role"], "content": m["content"]}
                    for m in st.session_state.messages
                ],
                temperature=temp_values,
                stream=True,
                max_tokens=max_tokens,
            )
            response = st.write_stream(stream)

        except Exception as e:
            response = "😵 Parece que algo deu errado!\
                        \n O modelo pode estar sendo atualizado ou há um problema no sistema.\
                        \n Tente novamente mais tarde."
            st.write(response)
            st.write("Esta foi a mensagem de erro:")
            st.write(e)

    st.session_state.messages.append({"role": "assistant", "content": response})