import os
import re
from pathlib import Path


def clean_text_in_files(input_dir, output_dir, file_extension=".txt"):
    """
    Очищает все файлы в директории, удаляя блоки think с помощью регулярного выражения.

    :param input_dir: Путь к исходной директории с файлами
    :param output_dir: Путь для сохранения очищенных файлов
    :param file_extension: Расширение обрабатываемых файлов
    """
    # Создаем выходную директорию, если её нет
    Path(output_dir).mkdir(parents=True, exist_ok=True)

    # Регулярное выражение для поиска и удаления блоков think
    pattern = re.compile(r"<think>.*?<\/think>", flags=re.DOTALL | re.IGNORECASE)

    # Счетчики для статистики
    processed_files = 0
    total_files = 0

    # Проходим по всем файлам в директории
    for filename in os.listdir(input_dir):
        if filename.endswith(file_extension):
            total_files += 1
            input_path = os.path.join(input_dir, filename)
            output_path = os.path.join(output_dir, filename)

            try:
                # Читаем файл
                with open(input_path, encoding="utf-8") as file:
                    content = file.read()

                # Применяем регулярное выражение
                cleaned_content = pattern.sub("", content)

                # Сохраняем результат
                with open(output_path, "w", encoding="utf-8") as file:
                    file.write(cleaned_content)

                processed_files += 1
                print(f"Обработан файл: {filename}")

            except Exception as e:
                print(f"Ошибка при обработке {filename}: {str(e)}")

    print(f"\nГотово! Обработано файлов: {processed_files}/{total_files}")


# Пример использования
if __name__ == "__main__":
    input_directory = "breed_descriptions"  # Исходная папка с файлами
    output_directory = "breed_descriptions"  # Папка для очищенных файлов

    clean_text_in_files(input_directory, output_directory)
