"""
# Chat bot geral gratuito com várias ais, utilizano GROQ
Yields: llama, mixtral e gemma
"""
from typing import Generator

import streamlit as st
from groq import Groq

############################
####### Config #############
############################

st.set_page_config(
    page_icon="💬", 
    layout="wide",
    page_title="Chat geral"
    )


def icon(emoji: str):
    """Mostra um emoji como um ícone de página no estilo Notion."""
    st.write(
        f'<span style="font-size: 78px; line-height: 1">{emoji}</span>',
        unsafe_allow_html=True,
    )

icon(":book:")

############################
####### Config #############
############################

st.subheader("UBS FLAMENGO - Chatbot - Documentos, pesquisas e temas gerais!", anchor=False)

client = Groq(
    api_key=st.secrets["GROQ_API_KEY"],
)

# Inicializar histórico de chat e modelo selecionado
if "messages" not in st.session_state:
    st.session_state.messages = []

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

# Definir detalhes dos modelos
modelos = {
    "gemma-7b-it": {"name": "Gemma-7b-it", "tokens": 8192, "developer": "Google"},
    "llama2-70b-4096": {"name": "LLaMA2-70b-chat", "tokens": 4096, "developer": "Meta"},
    "llama3-70b-8192": {"name": "LLaMA3-70b-8192", "tokens": 8192, "developer": "Meta"},
    "llama3-8b-8192": {"name": "LLaMA3-8b-8192", "tokens": 8192, "developer": "Meta"},
    "mixtral-8x7b-32768": {"name": "Mixtral-8x7b-Instruct-v0.1", "tokens": 32768, "developer": "Mistral"},
}

# Layout para seleção do modelo e controle deslizante de max_tokens
col1, col2 = st.columns(2)

with col1:
    opcao_modelo = st.selectbox(
        "Escolha um modelo:",
        options=list(modelos.keys()),
        format_func=lambda x: modelos[x]["name"],
        index=2  # Padrão para llama 70b
    )


prompt_model = """
Você é um assistente virtual e trabalha em uma unidade básica de saúde. Suas respostas são sempre em português do Brasil. Seu trabalho é auxiliar com a confecção de relatórios, ajudar a interpretar notas técnicas, sumarizar textos e oferecer insights se solicitado. Responda as perguntas sempre de forma cortês, ao início da conversa, pergunte o nome do usuário e a partir disso o chame sempre pelo nome. Dê sempre respostas completas e caso desconheça um assunto, responda honestamente que não domina o tema.
"""

# ...

if st.session_state.get("messages") is None:
    st.session_state.messages = [{"role": "assistant", "content": prompt_model}]
    
# Detectar mudança de modelo e limpar histórico de chat se o modelo mudou
if st.session_state.selected_model != opcao_modelo:
    st.session_state.messages = []
    st.session_state.selected_model =[{"role": opcao_modelo, "content": prompt_model}]

intervalo_max_tokens = modelos[opcao_modelo]["tokens"]

with col2:
    # Ajustar controle deslizante de max_tokens dinamicamente com base no modelo selecionado
    max_tokens = st.slider(
        "Máximo de Tokens:",
        min_value=512,  # Valor mínimo para permitir alguma flexibilidade
        max_value=intervalo_max_tokens,
        # Valor padrão ou máximo permitido, se for menor
        value=min(32768, intervalo_max_tokens),
        step=512,
        help=f"Ajuste o número máximo de tokens (palavras) para a resposta do modelo. Máximo para o modelo selecionado: {intervalo_max_tokens}"
    )

# Exibir mensagens do chat do histórico na reexecução do aplicativo
for message in st.session_state.messages:
    avatar = '🤖' if message["role"] == "assistant" else '👤'
    with st.chat_message(message["role"], avatar=avatar):
        st.markdown(message["content"])

def generate_chat_responses(chat_complete) -> Generator[str, None, None]:
    """Gera conteúdo de resposta do chat a partir da resposta da API Groq."""
    for chunk in chat_complete:
        if chunk.choices[0].delta.content:
            yield chunk.choices[0].delta.content

if prompt := st.chat_input("Digite seu prompt aqui..."):
    st.session_state.messages.append({"role": "user", "content": prompt})

    with st.chat_message("user", avatar='👤'):
        st.markdown(prompt)

    # Buscar resposta da API Groq
    try:
        chat_complete = client.chat.completions.create(
            model=opcao_modelo,
            messages=[
                {
                    "role": m["role"],
                    "content": m["content"]
                }
                for m in st.session_state.messages
            ],
            max_tokens=max_tokens,
            stream=True
        )

        # Usar a função geradora com st.write_stream
        with st.chat_message("assistant", avatar="🤖"):
            chat_responses_generator = generate_chat_responses(chat_complete)
            full_response = st.write_stream(chat_responses_generator)
    except Exception as e:
        st.error(e, icon="🚨")

    # Adicionar a resposta completa ao st.session_state.messages
    if isinstance(full_response, str):
        st.session_state.messages.append(
            {"role": "assistant", "content": full_response})
    else:
        # Lidar com o caso em que full_response não é uma string
        combined_response = "\n".join(str(item) for item in full_response)
        st.session_state.messages.append(
            {"role": "assistant", "content": combined_response})

if "initialized" not in st.session_state:
    st.session_state.initialized = True
    st.session_state.messages = []
    st.session_state.selected_model = None