# Importações necessárias
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import pandas as pd
import gradio as gr
from typing import List, Dict, Tuple, Optional

# Carregando o tokenizador e o modelo
tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium")

# Criando um dataframe com status de pedidos
dados_pedidos = {
    "numero_pedido": ["12345", "67890", "11121", "22232"],
    "status": ["Shipped", "Processing", "Delivered", "Cancelled"]
}
df_status_pedidos = pd.DataFrame(dados_pedidos)

# Função para verificar o status do pedido
def verificar_status_pedido(numero_pedido: str) -> str:
    try:
        status = df_status_pedidos[df_status_pedidos['numero_pedido'] == numero_pedido]['status'].iloc[0]
        return f'The status of your order {numero_pedido} is: {status}'
    except IndexError:
        return 'Order number not found. Please check and try again'

# Lista de palavras-chave para o status
palavras_chave_status = ['order', 'order status', 'status of my order', 'check my order', 'track my order', 'order update']

def responder(input_usuario: str, ids_historico_chat: Optional[torch.Tensor]) -> Tuple[str, torch.Tensor]:
    if any(keyword in input_usuario.lower() for keyword in palavras_chave_status):
        return 'Could you please enter your order number?', ids_historico_chat

    novo_usuario_input_ids = tokenizer.encode(input_usuario + tokenizer.eos_token, return_tensors='pt')
    
    if ids_historico_chat is not None:
        bot_input_ids = torch.cat([ids_historico_chat, novo_usuario_input_ids], dim=-1)
    else:
        bot_input_ids = novo_usuario_input_ids

    ids_historico_chat = model.generate(
        bot_input_ids,
        max_length=1000,
        pad_token_id=tokenizer.eos_token_id
    )
    resposta = tokenizer.decode(ids_historico_chat[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)
    return resposta, ids_historico_chat

# Criando interface
with gr.Blocks() as app:
    gr.Markdown("# Chat automatizado para E-Commerce")
    gr.Markdown("Envie uma pergunta sobre os seguintes pedidos em inglês: 12345, 67890, 11121, 22232")
    
    chatbot = gr.Chatbot(label="Chat")
    msg = gr.Textbox(label="Sua mensagem", placeholder='Digite sua mensagem aqui...')
    estado = gr.State(None)
    aguardando_numero_pedido = gr.State(False)

    def processar_entrada(input_usuario: str, historico: List[Dict[str, str]], 
                          ids_historico_chat: Optional[torch.Tensor], 
                          aguardando_numero_pedido: bool) -> Tuple[List[Dict[str, str]], Optional[torch.Tensor], bool, str]:
        if aguardando_numero_pedido:
            resposta = verificar_status_pedido(input_usuario)
            aguardando_numero_pedido = False
        else:
            resposta, ids_historico_chat = responder(input_usuario, ids_historico_chat)
        
        if resposta == 'Could you please enter your order number?':
            aguardando_numero_pedido = True
        
        historico.append((input_usuario, resposta))
        
        return historico, ids_historico_chat, aguardando_numero_pedido, ""

    msg.submit(
        processar_entrada,
        [msg, chatbot, estado, aguardando_numero_pedido],
        [chatbot, estado, aguardando_numero_pedido, msg]
    )

if __name__ == "__main__":
    app.launch(share=True)