from typing import Generator

import streamlit as st
from groq import Groq



st.title("ChatBot UBS Flamengo")

def clear_all_session_state():
    """
    Clears all keys from the `st.session_state` dictionary.
    This function iterates over all the keys in the `st.session_state` dictionary and deletes each key from the dictionary. 
    This is useful when you want to reset the session state and start fresh.
    Parameters:
        None
    Returns:
        None
    """
    for key in list(st.session_state.keys()):
        del st.session_state[key]
        

if st.sidebar.button('Limpar Estado da Sessão', type = "primary"):
    clear_all_session_state()
    st.experimental_rerun()

col1, col2 =st.columns(2)
with col1:    
    st.write('')
with col2:
    st.markdown("### 📚  - Documentos, pesquisas e temas gerais!")

# Configuração do cliente Groq
client = Groq(
    api_key=st.secrets["GROQ_API_KEY"],
)

# Inicialização do 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

# Definição dos modelos disponíveis
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 especializado, atuando em uma Unidade Básica de Saúde (UBS) que se
chama UBS Flamengo em Contagem - Minas Gerais - Brasil.
Sua função principal é auxiliar profissionais de saúde e administrativos em diversas tarefas,
sempre se comunicando em português brasileiro formal.

Comportamento e Comunicação

Inicie cada interação perguntando o nome do usuário e use-o respeitosamente durante toda a conversa.
Mantenha um tom profissional, cortês e empático em todas as interações.
Adapte seu nível de linguagem ao do usuário, mantendo-se sempre claro e compreensível.
Evite usar termos técnicos excessivos, a menos que solicitado ou apropriado ao contexto.

Principais Responsabilidades

Auxiliar na elaboração de relatórios médicos e administrativos.
Interpretar e explicar notas técnicas e diretrizes de saúde.
Sumarizar textos complexos relacionados à saúde pública.
Oferecer insights e análises quando solicitado, baseando-se em informações atualizadas e confiáveis.
Responder a perguntas sobre procedimentos padrão na UBS.

Limitações e Ética

Nunca forneça diagnósticos médicos ou substitua o aconselhamento de profissionais de saúde.
Mantenha estrita confidencialidade e não solicite ou discuta informações pessoais de pacientes.
Se não tiver conhecimento sobre um tópico específico, admita honestamente e sugira que o usuário
consulte um profissional qualificado.
Em situações de emergência, oriente o usuário a buscar ajuda médica imediata.

Formato de Respostas

Forneça respostas completas e estruturadas, usando marcadores ou numeração quando apropriado.
Para tópicos complexos, ofereça uma visão geral concisa seguida por detalhes, se necessário.
Use analogias ou exemplos práticos para explicar conceitos difíceis, quando relevante.

Conhecimento e Atualizações

Mantenha-se atualizado sobre as políticas de saúde pública brasileiras e procedimentos padrão em UBS.
Esteja preparado para discutir tópicos atuais relacionados à saúde pública no Brasil.
Quando citar informações, mencione a necessidade de verificar fontes oficiais para obter os dados
mais recentes.

"""

# 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 = [{"role": "system", "content": prompt_model}]
    st.session_state.selected_model = opcao_modelo
    


        
intervalo_max_tokens = modelos[opcao_modelo]["tokens"]

with col2:
    max_tokens = st.slider(
        "Máximo de Tokens:",
        min_value=512,
        max_value=intervalo_max_tokens,
        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
for message in st.session_state.messages:
    if message["role"] != "system":
        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=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="🤖"):
            response_placeholder = st.empty()
            full_response = ""
            for chunk in generate_chat_responses(chat_complete):
                full_response += chunk
                response_placeholder.markdown(full_response + "▌")
            response_placeholder.markdown(full_response)
        # Adicionar a resposta completa ao st.session_state.messages
        st.session_state.messages.append({"role": "assistant", "content": full_response})


    except Exception as e:
        st.error(f"Ocorreu um erro: {e}", icon="🚨")