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


dotenv.load_dotenv()
token = os.environ['HF_TOKEN']


# add models here
modelos = {
    'Felladrin/Llama-68M-Chat-v1': '[/INST]',
    'google/gemma-7b-it': '<start_of_turn>model\n',
    'google/gemma-2-2b-it': '<start_of_turn>model\n'
}

nome_modelo = st.selectbox('select a model:', options=modelos)
token_modelo = modelos[nome_modelo]

url = f"https://api-inference.huggingface.co/models/{nome_modelo}"
tokenizer = AutoTokenizer.from_pretrained(nome_modelo, token=token)

if 'modelo_atual' not in st.session_state or st.session_state['modelo_atual'] != nome_modelo:
    st.session_state['modelo_atual'] = nome_modelo
    st.session_state['mensagens'] = []
    
mensagens = st.session_state['mensagens']

area_chat = st.empty()
pergunta_usuario = st.chat_input('Do you question here: ')

if pergunta_usuario:
    mensagens.append({'role': 'user', 'content': pergunta_usuario})
    inputs = tokenizer.apply_chat_template(mensagens, tokenize=False, add_generation_prompt=True)
    json = {
        'inputs': inputs,
        'parameters': {'max_new_tokens': 1_000},
        'options': {'use_cache': False, 'wait_for_model': True},
    }
    headers = {
        'Authorization': f'Bearer {token}',
    }
    
    # Send request and verify response
    response = requests.post(url, json=json, headers=headers).json()
    
    # verify if response is valid
    if isinstance(response, list) and len(response) > 0 and 'generated_text' in response[0]:
        resposta_chatbot = response[0]['generated_text'].split(token_modelo)[-1]
        mensagens.append({'role': 'assistant', 'content': resposta_chatbot})
    else:
        mensagens.append({'role': 'assistant', 'content': "Desculpe, algo deu errado ao processar sua solicitação."})

with area_chat.container():
    for mensagem in mensagens:
        chat = st.chat_message(mensagem['role'])
        chat.markdown(mensagem['content'])