# app.py

import gradio as gr
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM # Используем Auto классы для большей гибкости
import random
import re
import torch # Убедимся, что torch импортирован

# --- Выбор Модели ---
# Заменяем на самую маленькую DeepSeek Coder Instruct
# !!! ПРЕДУПРЕЖДЕНИЕ: МОЖЕТ НЕ ЗАПУСТИТЬСЯ ИЛИ БЫТЬ ОЧЕНЬ МЕДЛЕННОЙ НА FREE CPU !!!
CODE_MODEL_NAME = "deepseek-ai/deepseek-coder-1.3b-instruct"

# --- Загрузка моделей ---
GENERATOR_LOADED = False
QA_LOADED = False
generator = None
qa_pipeline = None

# 1. Модель для кодирования/генерации
try:
    # Для моделей кода часто лучше использовать AutoModelForCausalLM и AutoTokenizer напрямую
    # Вместо pipeline('text-generation') это дает больше контроля
    print(f"Загрузка токенизатора: {CODE_MODEL_NAME}")
    tokenizer = AutoTokenizer.from_pretrained(CODE_MODEL_NAME, trust_remote_code=True)
    print(f"Загрузка модели: {CODE_MODEL_NAME}")
    # Загружаем с torch_dtype=torch.float16 для экономии памяти, если возможно
    try:
         generator_model = AutoModelForCausalLM.from_pretrained(
             CODE_MODEL_NAME,
             trust_remote_code=True,
             torch_dtype=torch.float16 # Пытаемся загрузить в половинной точности
             )
    except Exception as e_half: # Если float16 не поддерживается/вызывает ошибку, пробуем float32
        print(f"Не удалось загрузить в float16 ({e_half}), пробую float32...")
        generator_model = AutoModelForCausalLM.from_pretrained(
             CODE_MODEL_NAME,
             trust_remote_code=True
             )

    # Создаем pipeline вручную
    generator = pipeline('text-generation', model=generator_model, tokenizer=tokenizer, device=-1) # device=-1 значит CPU

    print(f"Генератор кода/текста ({CODE_MODEL_NAME}) загружен.")
    GENERATOR_LOADED = True
except Exception as e:
    print(f"ОШИБКА: Не удалось загрузить модель {CODE_MODEL_NAME}! {e}")
    generator = None
    generator_model = None
    tokenizer = None
    GENERATOR_LOADED = False

# 2. QA Модель (оставляем как есть или можно убрать, если не нужна)
QA_MODEL_NAME = 'timpal0l/mdeberta-v3-base-squad2'
try:
    qa_pipeline = pipeline('question-answering', model=QA_MODEL_NAME)
    print(f"QA модель ({QA_MODEL_NAME}) загружена.")
    QA_LOADED = True
    # !!! ВНИМАНИЕ: Две модели могут потребовать слишком много RAM !!!
    # !!! Возможно, стоит закомментировать загрузку QA, если возникают проблемы с памятью !!!
except Exception as e:
    print(f"ОШИБКА: Не удалось загрузить QA модель! {e}")
    qa_pipeline = None
    QA_LOADED = False


# --- Встроенные знания (можно сократить, т.к. модель сама умнее) ---
knowledge_base = {
    "кто ты": f"Я Nova (использую модель {CODE_MODEL_NAME}), работающая на Hugging Face Spaces. Я специализируюсь на помощи с кодом и ответами на вопросы по программированию.",
    "что ты умеешь": "Я могу пытаться генерировать код, объяснять его, находить ошибки (в определенной степени), отвечать на вопросы по программированию и поддерживать диалог.",
    "как дела": "Работаю над кодом! А если серьезно - функционирую нормально.",
    "помощь": "Спросите меня о программировании, попросите написать фрагмент кода (например, 'напиши функцию на python для суммирования списка'), или просто пообщаемся.",
    "python": "Python - отличный высокоуровневый язык, известный своей читаемостью. Хорошо подходит для веб-разработки, анализа данных, ИИ и автоматизации.",
    "javascript": "JavaScript - основной язык веб-фронтенда, позволяет делать страницы интерактивными. С Node.js используется и на бэкенде.",
    # Убрали много базовых вещей, т.к. модель должна знать их сама
}

# --- Утилиты (clean_text, clean_generated_output - как раньше) ---
def clean_text(text):
    if not isinstance(text, str): return ""
    text = text.lower().strip()
    text = re.sub(r"^[?!.,\s]+|[?!.,\s]+$", "", text)
    text = re.sub(r"\s+", " ", text)
    return text

def clean_generated_output(generated_text, prompt):
     # Простая очистка - убираем промпт если он есть в начале
    if not generated_text or not isinstance(generated_text, str): return ""
    if not prompt or not isinstance(prompt, str): prompt = ""
    cleaned = generated_text
    prompt_clean = prompt.strip()
    if prompt_clean and cleaned.startswith(prompt_clean):
        cleaned = cleaned[len(prompt_clean):].strip()

    # Попробуем убрать стандартные фразы, которые модель может добавить
    cleaned = re.sub(r'^\s*(ответ|бот|вот код|конечно|here is the code)\s*[:\-]?\s*', '', cleaned, flags=re.IGNORECASE).strip()

    return cleaned


# --- Основная Функция Обработки (Упрощенная, больше полагаемся на модель) ---
def respond(message, chat_history):
    print("-" * 30)
    print(f"ВХОД: '{message}'")
    user_message_clean = clean_text(message)
    bot_response = ""

    # 1. Проверка на точное совпадение в knowledge_base (ОСТАВЛЯЕМ для мета-вопросов)
    if user_message_clean in knowledge_base:
        bot_response = knowledge_base[user_message_clean]
        print(f"Ответ [Точный]: {bot_response[:100]}...")

    # 2. Основная генерация (используем модель DeepSeek Coder)
    if not bot_response and GENERATOR_LOADED and generator:
        print("Генерация ответа с помощью модели кода...")
        # Формируем промпт в формате, который модель ожидает (часто нужен специальный формат для Instruct-моделей)
        # Для DeepSeek Coder Instruct формат может быть таким:
        prompt_list = []
        prompt_list.append("### Instruction:")
        # Добавим контекст из чата (если есть)
        for user_msg, bot_msg in chat_history[-2:]: # Последние 2 обмена для контекста
             prompt_list.append(f"User: {user_msg}")
             if bot_msg: prompt_list.append(f"Assistant: {bot_msg}") # Используем Assistant
        prompt_list.append(f"{message}") # Текущее сообщение как инструкция/вопрос
        prompt_list.append("\n### Response:") # Приглашение для ответа

        full_prompt = "\n".join(prompt_list)
        print(f"Промпт для DeepSeek: ...{full_prompt[-600:]}") # Логируем конец

        try:
            # Параметры генерации для кода
            # `max_new_tokens` - ограничивает длину *сгенерированного* текста
            # `temperature` - контроль случайности (ниже = более предсказуемо)
            # `top_p`, `top_k` - другие методы семплирования
            # `eos_token_id` - ID токена конца последовательности (может помочь модели остановиться)
            generated_output = generator(
                full_prompt,
                max_new_tokens=250, # Ограничим генерацию
                temperature=0.7,
                top_k=50,
                top_p=0.95,
                # eos_token_id=tokenizer.eos_token_id if tokenizer else None, # Помогает модели остановиться
                pad_token_id=tokenizer.eos_token_id if tokenizer else 50256, # Часто нужно указать
                num_return_sequences=1,
                do_sample=True
                )[0]['generated_text']

            print(f"Сырой ответ DeepSeek: {generated_output}")

            # Очистка - убираем весь промпт
            bot_response = clean_generated_output(generated_output, full_prompt)
            print(f"Очищенный ответ DeepSeek: {bot_response}")

            # Дополнительно убираем инструкции, которые модель могла повторить
            bot_response = bot_response.replace("### Instruction:", "").replace("### Response:", "").strip()

            if not bot_response or len(bot_response) < 5:
                bot_response = knowledge_base.get("не найдено", "Извините, не смог сгенерировать ответ.")
                print("Генерация DeepSeek не удалась или слишком короткая.")

        except Exception as e:
            print(f"Ошибка при генерации DeepSeek: {e}")
            bot_response = knowledge_base.get("ошибка генерации", f"Произошла ошибка при генерации: {e}")

    # 3. Ответ по умолчанию
    if not bot_response:
        # Если генератор не загружен или ничего не получилось
        if not GENERATOR_LOADED:
            bot_response = "Модель для генерации кода не загружена."
        else:
            bot_response = knowledge_base.get("не найдено", "К сожалению, я пока не знаю, как на это ответить.")
        print("Ответ [Заглушка/Ошибка]")


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


# --- Gradio Интерфейс (CSS тот же) ---
custom_css = """
.gradio-container { background-color: #f4f7f9; border-radius: 15px; padding: 25px; color: #333; }
h1 { color: #1a237e; text-align: center; font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif; margin-bottom: 15px; font-weight: 600; }
#title-markdown p { text-align: center; color: #546e7a; margin-top: -10px; margin-bottom: 25px;}
#chatbot { background-color: #eceff1; border-radius: 12px; border: 1px solid #cfd8dc; padding: 10px; box-shadow: inset 0 1px 3px rgba(0,0,0,0.05); }
#chatbot .user-message .message-bubble-border { border: none !important; }
#chatbot .user-message .message-bubble { background: #007bff !important; color: white !important; border-radius: 18px 18px 5px 18px !important; padding: 12px 18px !important; margin: 8px 5px 8px 0 !important; align-self: flex-end !important; max-width: 80% !important; box-shadow: 0 2px 4px rgba(0, 123, 255, 0.2); word-wrap: break-word; text-align: left;}
#chatbot .bot-message .message-bubble-border { border: none !important; }
#chatbot .bot-message .message-bubble { background: #ffffff !important; color: #212529 !important; border: 1px solid #dee2e6 !important; border-radius: 18px 18px 18px 5px !important; padding: 12px 18px !important; margin: 8px 0 8px 5px !important; align-self: flex-start !important; max-width: 80% !important; box-shadow: 0 2px 4px rgba(0, 0, 0, 0.05); word-wrap: break-word; text-align: left;}
#chatbot .bot-message .message-bubble pre { background-color: #f8f9fa; border: 1px solid #e9ecef; border-radius: 6px; padding: 10px; margin: 8px 0 5px 0; overflow-x: auto; word-wrap: normal; }
#chatbot .bot-message .message-bubble pre code { background-color: transparent !important; color: #2d3748; font-family: 'Consolas', 'Monaco', 'Courier New', monospace; font-size: 0.9em; padding: 0; white-space: pre; }
textarea { border: 1px solid #ced4da !important; border-radius: 10px !important; padding: 12px 15px !important; background-color: #ffffff; transition: border-color 0.3s ease, box-shadow 0.3s ease; font-size: 1rem; }
textarea:focus { border-color: #80bdff !important; box-shadow: 0 0 0 0.2rem rgba(0, 123, 255, 0.25); outline: none; }
button { border-radius: 10px !important; padding: 10px 15px !important; transition: background-color 0.2s ease, transform 0.1s ease !important; font-weight: 500 !important; border: none !important; }
button:active { transform: translateY(1px); }
button.primary { background-color: #007bff !important; color: white !important; }
button.primary:hover { background-color: #0056b3 !important; box-shadow: 0 2px 5px rgba(0, 123, 255, 0.3); }
button.secondary { background-color: #6c757d !important; color: white !important; }
button.secondary:hover { background-color: #5a6268 !important; box-shadow: 0 2px 5px rgba(108, 117, 125, 0.3); }
"""

with gr.Blocks(css=custom_css, theme=gr.themes.Soft(primary_hue=gr.themes.colors.blue, secondary_hue=gr.themes.colors.slate)) as demo:
    gr.Markdown("# 🤖 Nova (DeepSeek Test) 💡 <br> Используем модель для кода. Ответы могут быть ОЧЕНЬ медленными!")

    chatbot = gr.Chatbot(
        label="Диалог", height=550, elem_id="chatbot",
        show_copy_button=True, bubble_full_width=False,
        avatar_images=(None, "https://img.icons8.com/plasticine/100/bot.png")
        )
    with gr.Row():
        msg = gr.Textbox(
            label="Ваше сообщение", placeholder="Напиши функцию python для...",
            scale=4, show_label=False, container=False
        )
        submit_btn = gr.Button("➤ Отправить", variant="primary", scale=1, min_width=120)
        clear_btn = gr.Button("🗑️ Очистить", variant="secondary", scale=1, min_width=120)

    msg.submit(respond, [msg, chatbot], [msg, chatbot])
    submit_btn.click(respond, [msg, chatbot], [msg, chatbot])
    clear_btn.click(lambda: (None, []), None, [msg, chatbot], queue=False)

# Запуск
demo.queue()
demo.launch(debug=True)