# === ЗАГРУЗКА МОДЕЛИ (ИСПРАВЛЕННАЯ ЧАСТЬ) ===
print("🚀 Загрузка модели...")
model_id = "OpenRussianAI/OpenAirAI-X"
tokenizer = AutoTokenizer.from_pretrained(model_id)

# ВАЖНО: Убеждаемся, что у токенизатора есть pad_token
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

model = AutoModelForCausalLM.from_pretrained(model_id)
device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device).eval()
print(f"✅ Модель загружена на {device}")

# === ГЕНЕРАЦИЯ ОТВЕТА (ПОЛНОСТЬЮ ПЕРЕПИСАННАЯ) ===
def generate_response(message, history, username, current_chat_id):
    if not username:
        gr.Warning("Сначала введите имя пользователя!")
        return history, gr.update(), current_chat_id, gr.update()
    
    if not message.strip():
        return history, gr.update(), current_chat_id, gr.update()
    
    # Добавляем сообщение пользователя в историю
    history = history + [{"role": "user", "content": message}]
    
    # Формируем промпт
    prompt = ""
    for msg in history:
        if msg["role"] == "user":
            prompt += f"Пользователь: {msg['content']}\n"
        elif msg["role"] == "assistant":
            prompt += f"AI: {msg['content']}\n"
    prompt += "AI:"
    
    inputs = tokenizer(prompt, return_tensors="pt").to(device)
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=128, # Уменьшим, чтобы модель не уходила в бред
            temperature=0.7,
            top_p=0.9,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id,
            eos_token_id=tokenizer.eos_token_id, # Важно!
            repetition_penalty=1.2 # Штраф за повторения, чтобы убрать цикл "Я OpenAI"
        )
    
    # Декодируем только НОВУЮ часть текста
    generated_ids = outputs[0][inputs.input_ids.shape[-1]:]
    response_text = tokenizer.decode(generated_ids, skip_special_tokens=True)
    
    # Очистка от мусора и обрезка
    # Обычно модель может начать ответ с пробела или новой строки
    ai_response = response_text.strip()
    
    # Если модель все же написала "Пользователь:" или "AI:" в конце, обрезаем
    stop_words = ["Пользователь:", "User:", "\n\n"]
    for stop_word in stop_words:
        if stop_word in ai_response:
            ai_response = ai_response.split(stop_word)[0].strip()
            
    # Если ответ пустой после очистки
    if not ai_response:
        ai_response = "..."

    history = history + [{"role": "assistant", "content": ai_response}]
    
    if username and current_chat_id:
        history_data = load_history(username)
        if current_chat_id not in history_data:
            history_data[current_chat_id] = {
                "title": message[:40] + ("..." if len(message) > 40 else ""),
                "created": datetime.now().isoformat(),
                "messages": []
            }
        history_data[current_chat_id]["messages"] = history
        save_history(username, history_data)
    
    chat_list_choices = get_chat_list(username)
    return history, gr.update(value=""), current_chat_id, gr.update(choices=chat_list_choices)