import os
import logging
from transformers import AutoModelForCausalLM, AutoTokenizer
from telegram import Update, ReplyKeyboardRemove
from telegram.ext import (
    ApplicationBuilder,
    CommandHandler,
    MessageHandler,
    filters,
    ConversationHandler,
)
import re  # Для поиска арифметических выражений
import sympy  # Для решения математических задач

# Настройка логирования
logging.basicConfig(
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    level=logging.INFO
)

# Загрузка модели и токенизатора CodeGen
MODEL_NAME = "Salesforce/codegen-350M-multi"  # Уменьшенная версия модели
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)

# Устанавливаем pad_token, чтобы он отличался от eos_token
tokenizer.pad_token = tokenizer.eos_token

# Глобальные переменные для хранения контекста
chat_history = {}
MAX_HISTORY_LENGTH = 10  # Ограничиваем длину истории диалога

# Состояния для ConversationHandler
CHAT = range(1)

# Функция для генерации ответа с использованием CodeGen
def generate_response(prompt, max_new_tokens=1000, temperature=0.5):
    try:
        # Токенизация входных данных
        inputs = tokenizer(prompt, return_tensors='pt', padding=True, truncation=True)

        # Генерация ответа
        output_ids = model.generate(
            inputs.input_ids,
            attention_mask=inputs.attention_mask,  # Передаем attention_mask
            max_new_tokens=max_new_tokens,  # Используем max_new_tokens вместо max_length
            pad_token_id=tokenizer.eos_token_id,
            no_repeat_ngram_size=3,
            do_sample=True,
            top_k=50,
            top_p=0.95,
            temperature=temperature
        )

        # Декодирование ответа
        response = tokenizer.decode(output_ids[0], skip_special_tokens=True)
        return response
    except Exception as e:
        logging.error(f"Ошибка при генерации ответа: {e}")
        return "Произошла ошибка при обработке вашего запроса."

# Функция для решения математических задач
def solve_math_problem(problem):
    try:
        # Удаляем лишние символы и пробелы
        problem = re.sub(r'\s+', '', problem)

        # Используем sympy для решения задачи
        result = sympy.sympify(problem)
        return f"Результат: {result}"
    except Exception as e:
        logging.error(f"Ошибка при решении задачи: {e}")
        return "Не удалось решить задачу. Убедитесь, что вы ввели корректное выражение."

# Обработчик команды /start
async def start(update: Update, context):
    user_id = update.message.from_user.id
    chat_history[user_id] = []  # Инициализация контекста для пользователя
    await update.message.reply_text(
        'Привет! Я ваш Telegram-бот с интеграцией CodeGen. Я могу генерировать готовый код, решать примеры, отвечать на сложные вопросы и вести диалог. Задавайте вопросы!',
        reply_markup=ReplyKeyboardRemove()
    )
    return CHAT

# Обработчик сообщений
async def handle_message(update: Update, context):
    user_id = update.message.from_user.id
    user_input = update.message.text

    # Добавляем сообщение в историю диалога
    add_to_history(user_id, f"Пользователь: {user_input}")

    # Проверяем, содержит ли сообщение математическое выражение
    if re.search(r'\d+[\+\-\*\/]\d+', user_input):
        # Если да, решаем пример
        response = solve_math_problem(user_input)
    elif "напиши код" in user_input.lower() or "напиши программу" in user_input.lower():
        # Если пользователь просит написать код, генерируем готовый код
        response = generate_response(user_input, max_new_tokens=2000, temperature=0.5)
    else:
        # Если нет, генерируем ответ на сложный вопрос
        full_context = "\n".join(chat_history[user_id])
        response = generate_response(full_context, max_new_tokens=1000, temperature=0.5)

    # Добавляем ответ в историю диалога
    add_to_history(user_id, f"Бот: {response}")

    # Отправка ответа пользователю
    await update.message.reply_text(response)
    return CHAT

# Функция для добавления сообщения в историю диалога
def add_to_history(user_id, message):
    if user_id not in chat_history:
        chat_history[user_id] = []
    chat_history[user_id].append(message)
    if len(chat_history[user_id]) > MAX_HISTORY_LENGTH:
        chat_history[user_id].pop(0)  # Удаляем самое старое сообщение

# Основная функция
if __name__ == '__main__':
    # Создание приложения с увеличенным временем ожидания
    application = ApplicationBuilder().token(os.getenv("TOKEN")).connect_timeout(60).read_timeout(60).build()

    # Добавление обработчиков
    conversation_handler = ConversationHandler(
        entry_points=[CommandHandler("start", start)],
        states={
            CHAT: [MessageHandler(filters.TEXT & ~filters.COMMAND, handle_message)],
        },
        fallbacks=[],
    )
    application.add_handler(conversation_handler)

    # Запуск бота
    application.run_polling()