import gradio as gr
import os
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
import torch

# Инициализация модели для русского языка
model_name = "microsoft/DialoGPT-medium"  # Можно заменить на sberbank-ai/rugpt3small_based_on_gpt2

# Загрузка модели и токенизатора
try:
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForCausalLM.from_pretrained(model_name)
    tokenizer.pad_token = tokenizer.eos_token
except Exception as e:
    print(f"Ошибка загрузки модели: {e}")
    # Fallback на более простую модель
    model_name = "gpt2"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForCausalLM.from_pretrained(model_name)

# Системный промпт для образовательного чат-бота
SYSTEM_PROMPT = """
Ты - образовательный помощник преподавателя университета. Твоя задача:
- Помогать студентам с учебными вопросами
- Объяснять сложные концепции простым языком
- Давать конструктивные советы по обучению
- Быть терпеливым и поддерживающим
- Отвечать на русском языке

Всегда:
- Будь дружелюбным и профессиональным
- Объясняй пошагово
- Приводи примеры
- Задавай уточняющие вопросы если нужно
"""

def format_prompt(message, history):
    """Форматирование промпта с учетом истории"""
    prompt = SYSTEM_PROMPT + "\n\n"
    
    # Добавляем историю разговора
    for user_msg, bot_msg in history:
        prompt += f"Студент: {user_msg}\nПреподаватель: {bot_msg}\n\n"
    
    prompt += f"Студент: {message}\nПреподаватель:"
    return prompt

def generate_response(message, history, temperature=0.7, max_length=200):
    """Генерация ответа от модели"""
    try:
        # Форматируем промпт
        prompt = format_prompt(message, history)
        
        # Токенизируем
        inputs = tokenizer.encode(prompt, return_tensors="pt", max_length=1024, truncation=True)
        
        # Генерируем ответ
        with torch.no_grad():
            outputs = model.generate(
                inputs,
                max_length=inputs.shape[1] + max_length,
                temperature=temperature,
                do_sample=True,
                pad_token_id=tokenizer.eos_token_id,
                no_repeat_ngram_size=2,
                top_p=0.9
            )
        
        # Декодируем ответ
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        # Извлекаем только новую часть ответа
        response = response[len(prompt):].strip()
        
        # Очищаем ответ от лишних символов
        response = response.split("Студент:")[0].strip()
        
        if not response:
            response = "Извините, не могу сформулировать ответ. Попробуйте переформулировать вопрос."
            
        return response
        
    except Exception as e:
        return f"Произошла ошибка: {str(e)}. Попробуйте еще раз."

def chat_function(message, history, temperature, max_length):
    """Основная функция чата"""
    if not message.strip():
        return history, ""
    
    # Генерируем ответ
    response = generate_response(message, history, temperature, max_length)
    
    # Добавляем в историю
    history.append((message, response))
    
    return history, ""

# Создание интерфейса Gradio
with gr.Blocks(
    title="Образовательный чат-бот",
    theme=gr.themes.Soft(primary_hue="blue"),
    css="""
    .gradio-container {
        max-width: 800px !important;
        margin: auto !important;
    }
    """
) as demo:
    
    gr.HTML("""
    <div style="text-align: center; margin-bottom: 20px;">
        <h1>📚 Образовательный чат-бот</h1>
        <p>Ваш персональный помощник для обучения</p>
    </div>
    """)
    
    with gr.Row():
        with gr.Column(scale=3):
            chatbot = gr.Chatbot(
                height=500,
                placeholder="Привет! Я ваш образовательный помощник. Задайте мне любой вопрос по учебе!",
                show_label=False
            )
            
            with gr.Row():
                message = gr.Textbox(
                    placeholder="Введите ваш вопрос...",
                    show_label=False,
                    lines=2,
                    scale=4
                )
                send_btn = gr.Button("Отправить", variant="primary", scale=1)
        
        with gr.Column(scale=1):
            gr.HTML("<h3>⚙️ Настройки</h3>")
            
            temperature = gr.Slider(
                minimum=0.1,
                maximum=1.0,
                value=0.7,
                step=0.1,
                label="Креативность",
                info="Выше = более креативные ответы"
            )
            
            max_length = gr.Slider(
                minimum=50,
                maximum=300,
                value=200,
                step=50,
                label="Длина ответа",
                info="Максимальная длина ответа"
            )
            
            clear_btn = gr.Button("🗑️ Очистить чат", variant="secondary")
            
            gr.HTML("""
            <div style="margin-top: 20px; padding: 10px; background-color: #f0f0f0; border-radius: 5px;">
                <h4>💡 Примеры вопросов:</h4>
                <ul>
                    <li>Объясни теорему Пифагора</li>
                    <li>Как решать квадратные уравнения?</li>
                    <li>Что такое фотосинтез?</li>
                    <li>Помоги с написанием эссе</li>
                </ul>
            </div>
            """)
    
    # Обработчики событий
    send_btn.click(
        chat_function,
        inputs=[message, chatbot, temperature, max_length],
        outputs=[chatbot, message]
    )
    
    message.submit(
        chat_function,
        inputs=[message, chatbot, temperature, max_length],
        outputs=[chatbot, message]
    )
    
    clear_btn.click(
        lambda: ([], ""),
        outputs=[chatbot, message]
    )

# Запуск приложения
if __name__ == "__main__":
    demo.launch(
        share=False,
        server_name="0.0.0.0",
        server_port=7860
    )