import os
import requests
import dotenv
import streamlit as st
from transformers import AutoTokenizer

# Função para carregar as configurações iniciais
def carregar_configuracoes():
    dotenv.load_dotenv()
    return os.environ['TOKEN_HF']

# Função para inicializar o modelo e o tokenizador
def inicializar_modelo_e_tokenizador(modelo, token):
    tokenizer = AutoTokenizer.from_pretrained(modelo, token=token)
    return tokenizer

# Função para configurar a sessão do usuário
def configurar_sessao(modelo):
    if 'modelo_atual' not in st.session_state or st.session_state['modelo_atual'] != modelo:
        st.session_state['modelo_atual'] = modelo
        st.session_state['mensagens'] = [
            {"role": "system", "content": "Você é um bot para investimentos, treinado para ajudar os usuários a gerenciar e otimizar seus portfólios financeiros."
                                          "Seu objetivo é fornecer análises de mercado, sugestões de investimentos personalizadas e educar os usuários sobre estratégias financeiras."
                                          "Você pode responder a perguntas sobre ações, fundos, criptomoedas, imóveis e outros tipos de investimentos."
                                          "Além disso, você oferece insights sobre tendências de mercado e ajuda os usuários a tomar decisões informadas com base no perfil de risco deles."
                                          "Seja sucinto em suas respostas, resumindo-as para no máximo 200 caracteres."},
            {"role": "assistant", "content": "Olá, seja bem vindo ao assistente financeiro, como posso te ajudar?"}
        ]
        st.markdown("""
            <style>
            .reportview-container .markdown-text-container p {
                line-height: 1.5;
                word-wrap: break-word;
            }
            </style>
            """, unsafe_allow_html=True)

# Função para classificar a entrada do usuário
def classificar_entrada(pergunta_usuario, opcoes_classificacao, headers, url_classificacao, limite_score=70):
    payload = {
        "inputs": pergunta_usuario,
        "parameters": {
            "candidate_labels": opcoes_classificacao
        }
    }
    response = requests.post(url_classificacao, headers=headers, json=payload)
    
    if response.status_code == 200:
        result = response.json()
        score = result['scores'][0] * 100
        label = result['labels'][0]
        
        if score > limite_score:
            return label, score
        else:
            print(f'Não atingiu mais que {limite_score}%. Score máximo {score:.2f}% para "{label}"')
            return None, score
    else:
        print(f"Erro na Classificação: {response.status_code}")
        print(response.json())
        return None, None

# Função para gerar uma resposta com base na classificação
def gerar_resposta(label, opcoes_classificacao):
    if label == opcoes_classificacao[0]:
        novo_pront = (
            "Extrato de investimentos.\n"
			"- Você possui os seguintes investimentos: \n"
			"- Valor em ações: 3.800,00 \n"
			"- Valor em Criptomoedas: 10.200,00 \n"
            "- CC: 5.000,00 \n"
			"- CP: 3.200 \n"
        )
    else:
        novo_pront = (
            "Esse é seu assessor de investimentos.\n"
			"- Você possui o assessor: \n"
			"- João da Silva \n"
			"- Tel: (11) 9 1111-1111 \n"
			"- e-mail: joao@investimentos.com \n"
			"- Empresa: Investimentos-SA \n"
        )
    
    print(novo_pront)
    return novo_pront

# Função para processar a conversa e exibir no Streamlit
def processar_conversa(url, headers, tokenizer, mensagens, token_modelo):
    template = tokenizer.apply_chat_template(mensagens, tokenize=False, add_generation_prompt=True)
    json = {
        'inputs': template,
        'parameters': {'max_new_tokens': 600},
        'options': {'use_cache': False, 'wait_for_model': True},
    }
    response = requests.post(url, json=json, headers=headers).json()
    print(response)
    mensagem_chatbot = response[0]['generated_text'].split(token_modelo)[-1]
    mensagens.append({'role': 'assistant', 'content': mensagem_chatbot})
    return mensagens

# Função para exibir as mensagens no chat
def exibir_mensagens(area_chat, mensagens):
    with area_chat.container():
        for mensagem in mensagens:
            if mensagem['role'] != 'system':  # Filtra as mensagens do tipo 'system'
                chat = st.chat_message(mensagem['role'])
                chat.markdown(mensagem['content'])
                

# Função principal que orquestra as chamadas
def main():
    token = carregar_configuracoes()
    headers = {'Authorization': f'Bearer {token}'}
    
    opcoes_classificacao =  ["Ivestimentos", "assessor"]
    modelo = 'meta-llama/Meta-Llama-3.1-8B-Instruct'
    token_modelo = '<|end_header_id|>\n'  # Adicionando a variável token_modelo
    
    configurar_sessao(modelo)
    
    nome_modelo = st.session_state['modelo_atual'] 
    tokenizer = inicializar_modelo_e_tokenizador(nome_modelo, token)
    url = f'/static-proxy?url=https%3A%2F%2Fapi-inference.huggingface.co%2Fmodels%2F%7Bnome_modelo%7D' 
    url_classificacao = f'/static-proxy?url=https%3A%2F%2Fapi-inference.huggingface.co%2Fmodels%2Ffacebook%2Fbart-large-mnli'
    
    mensagens = st.session_state['mensagens']
    area_chat = st.empty()
    exibir_mensagens(area_chat, mensagens)  # Exibe as mensagens iniciais, incluindo a mensagem de boas-vindas
    
    pergunta_usuario = st.chat_input('Faça sua pergunta aqui: ')
    
    if pergunta_usuario:
        label, score = classificar_entrada(pergunta_usuario, opcoes_classificacao, headers, url_classificacao)
        
        if label:
            novo_pront = gerar_resposta(label, opcoes_classificacao)
            mensagens.append({"role": "system", "content": novo_pront})
        
        mensagens.append({'role': 'user', 'content': pergunta_usuario})
        mensagens = processar_conversa(url, headers, tokenizer, mensagens, token_modelo)
        exibir_mensagens(area_chat, mensagens)
    
    print(mensagens)

# Executa a função principal
if __name__ == "__main__":
    main()