# app.py

import gradio as gr
from transformers import pipeline, set_seed
import random

# Загружаем модель для генерации текста (text-generation pipeline)
# Выберите модель:
# 'distilgpt2' - маленькая английская модель GPT-2
# 'sberbank-ai/rugpt3small_based_on_gpt2' - маленькая русская модель (может требовать больше ресурсов)
# 'bigscience/bloomz-560m' - многоязычная модель поменьше (может быть медленной на CPU)

# Попробуем с distilgpt2 для начала, т.к. она легче всего запустится на CPU
# Если хотите русскую, раскомментируйте строку с 'sberbank-ai...' и закомментируйте distilgpt2
# generator = pipeline('text-generation', model='distilgpt2', max_new_tokens=50) # Ограничим длину генерации
# generator = pipeline('text-generation', model='sberbank-ai/rugpt3small_based_on_gpt2', max_new_tokens=50)
generator = pipeline('text-generation', model='distilgpt2', max_length=70, num_return_sequences=1, truncation=True) # Альтернативные параметры

print("Модель загружена!") # Для отладки в логах

# Установка зерна для воспроизводимости (опционально)
# set_seed(random.randint(0, 10000)) # Используем случайное зерно для разнообразия

# Функция, которая будет обрабатывать ввод пользователя и историю чата
def respond(message, chat_history):
    """
    Принимает сообщение пользователя и историю чата, возвращает ответ бота.
    """
    print(f"Получено сообщение: {message}")
    print(f"История чата: {chat_history}")

    # Формируем промпт для модели (можно экспериментировать)
    # Вариант 1: Просто сообщение пользователя
    prompt = message

    # Вариант 2: Добавляем немного истории (простой)
    # history_text = "\n".join([f"User: {m}\nBot: {b}" for m, b in chat_history[-2:]]) # Последние 2 обмена
    # prompt = f"{history_text}\nUser: {message}\nBot:"

    try:
        # Генерируем ответ
        # Обратите внимание: 'text-generation' часто лучше работает, если ей дать начало фразы
        # Попробуем сделать так, чтобы она просто продолжила мысль или ответила.
        # Можно добавить "Ответ:", "Бот:", и т.п. в конец prompt, если модель так лучше реагирует.

        bot_response = generator(prompt)[0]['generated_text']

        # Очистка ответа: убираем исходный промпт из сгенерированного текста
        # (pipeline часто возвращает промпт + генерацию)
        if bot_response.startswith(prompt):
            bot_response = bot_response[len(prompt):].strip()

        # Иногда модели могут генерировать пустые ответы или спецсимволы, добавим проверку
        if not bot_response or len(bot_response.strip()) == 0:
            bot_response = "... (модель не сгенерировала ответ)"

        print(f"Сгенерированный ответ: {bot_response}")

    except Exception as e:
        print(f"Ошибка при генерации: {e}")
        bot_response = f"Ой, произошла ошибка при генерации ответа: {e}"

    # Добавляем пару (сообщение пользователя, ответ бота) в историю
    chat_history.append((message, bot_response))

    # Возвращаем пустую строку для очистки поля ввода и обновленную историю
    return "", chat_history

# --- Создание интерфейса Gradio ---
with gr.Blocks() as demo:
    gr.Markdown("# Простой Чат-Бот на Hugging Face")
    gr.Markdown("Введите ваше сообщение и нажмите Enter.")

    # Компонент чат-бота
    chatbot = gr.Chatbot(label="Диалог")

    # Поле ввода текста
    msg = gr.Textbox(label="Ваше сообщение", placeholder="Напишите что-нибудь...")

    # Кнопка очистки (опционально)
    clear = gr.Button("Очистить диалог")

    # Связываем компоненты:
    # Когда пользователь отправляет сообщение (нажатием Enter в Textbox или Button, если бы она была)
    # вызывается функция respond с текущим сообщением и историей чата.
    # Результат функции (пустая строка и новая история) обновляет Textbox (msg) и Chatbot (chatbot).
    msg.submit(respond, [msg, chatbot], [msg, chatbot])

    # Связываем кнопку очистки
    clear.click(lambda: (None, None), None, [msg, chatbot], queue=False)

# Запускаем демо
# concurrency_count - сколько запросов может обрабатывать одновременно
# share=True - если хотите получить временную публичную ссылку (не нужно для Spaces)
demo.launch(debug=True) # Включаем debug для логов в интерфейсе Spaces
# demo.launch(server_name="0.0.0.0", server_port=7860) # Для запуска на некоторых платформах