import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import time

# Загрузка модели и токенизатора
MODEL_NAME = "microsoft/DialoGPT-medium"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)

# Инициализация NPC
npc_names = ["Agent Alpha", "Agent Beta", "Agent Gamma"]
chat_history = []

def generate_response(input_text):
    """
    Генерирует ответ от NPC на основе входного текста.
    """
    # Токенизация входного текста
    input_ids = tokenizer.encode(input_text + tokenizer.eos_token, return_tensors="pt")
    
    # Генерация ответа
    response_ids = model.generate(
        input_ids,
        max_length=1000,
        pad_token_id=tokenizer.eos_token_id,
        no_repeat_ngram_size=2,
        top_p=0.95,
        top_k=50
    )
    
    # Декодирование ответа
    response = tokenizer.decode(response_ids[:, input_ids.shape[-1]:][0], skip_special_tokens=True)
    return response

def npc_chat():
    """
    Обрабатывает диалог между NPC.
    """
    global chat_history
    
    # Начальное сообщение
    if not chat_history:
        chat_history.append(f"{npc_names[0]}: Welcome to Earth Web Space Networks!")
    
    for i, npc in enumerate(npc_names):
        # Получаем последнее сообщение из истории
        last_message = chat_history[-1].split(": ", 1)[1]
        
        # Генерируем ответ
        response = generate_response(last_message)
        new_message = f"{npc}: {response}"
        chat_history.append(new_message)
        
        # Ограничение длины истории чата
        if len(chat_history) > 20:
            chat_history.pop(0)
    
    # Возвращаем текущую историю чата
    formatted_chat = "\n".join(chat_history)
    return formatted_chat

# Gradio интерфейс
with gr.Blocks() as demo:
    gr.Markdown("# Earth Web Space Networks Live Chat")
    
    with gr.Row():
        chat_output = gr.Textbox(label="Chat History", lines=10, interactive=False)
    
    # Кнопка для обновления чата
    update_button = gr.Button("Update Chat")
    update_button.click(npc_chat, inputs=None, outputs=chat_output)

demo.launch()