# src/universal_preprocessor.py
"""
Универсальный модуль предобработки текста.
Обеспечивает стандартизацию пунктуации, замену специальных токенов
и обработку сокращений для приведения текста к единому стандарту.
"""

import re
from typing import Dict, List, Optional, Tuple
from dataclasses import dataclass


@dataclass
class PreprocessingConfig:
    """Конфигурация для предобработки текста."""
    replace_urls: bool = True
    replace_emails: bool = True
    replace_numbers: bool = True
    expand_abbreviations: bool = True
    normalize_punctuation: bool = True
    normalize_quotes: bool = True
    normalize_dashes: bool = True
    normalize_spaces: bool = True


# Регулярные выражения для поиска специальных элементов
RE_URL = re.compile(r'https?://\S+|www\.\S+', flags=re.I)
RE_EMAIL = re.compile(r'[\w.+-]+@[\w-]+\.[\w.-]+', flags=re.I)
RE_PHONE = re.compile(r'\+?[78][\s\-]?\(?\d{3}\)?[\s\-]?\d{3}[\s\-]?\d{2}[\s\-]?\d{2}')
RE_NUM = re.compile(r'(?<!\w)[+-]?\d[\d\.,]*')
RE_CURRENCY = re.compile(r'\d+[\s]*(?:руб|рублей|долл|долларов|евро|€|\$|₽)')
RE_PERCENT = re.compile(r'\d+[\s]*%')
RE_DATE = re.compile(r'\d{1,2}[./]\d{1,2}[./]\d{2,4}|\d{1,2}\s+(?:января|февраля|марта|апреля|мая|июня|июля|августа|сентября|октября|ноября|декабря)\s+\d{4}')

# Словарь сокращений для русского языка
COMMON_ABBREVIATIONS = {
    # Общие сокращения
    r'\bт\.е\.': 'то есть',
    r'\bт\.д\.': 'так далее',
    r'\bт\.п\.': 'тому подобное',
    r'\bи\.т\.д\.': 'и так далее',
    r'\bи\.т\.п\.': 'и тому подобное',
    r'\bт\.к\.': 'так как',
    r'\bт\.о\.': 'то есть',
    r'\bт\.н\.': 'так называемый',
    r'\bт\.с\.': 'то есть',
    r'\bт\.ч\.': 'то есть',
    
    # Временные сокращения
    r'\bг\.': 'год',
    r'\bгг\.': 'годы',
    r'\bв\.': 'век',
    r'\bвв\.': 'века',
    r'\bмин\.': 'минута',
    r'\bмин\.': 'минуты',
    r'\bсек\.': 'секунда',
    r'\bсек\.': 'секунды',
    r'\bчас\.': 'час',
    r'\bчасы\.': 'часы',
    
    # Географические сокращения
    r'\bул\.': 'улица',
    r'\bпр\.': 'проспект',
    r'\bпер\.': 'переулок',
    r'\bпл\.': 'площадь',
    r'\bнаб\.': 'набережная',
    r'\bш\.': 'шоссе',
    r'\bобл\.': 'область',
    r'\bр-н': 'район',
    r'\bг\.': 'город',
    r'\bс\.': 'село',
    r'\bд\.': 'деревня',
    r'\bп\.': 'поселок',
    
    # Организационные сокращения
    r'\bООО': 'общество с ограниченной ответственностью',
    r'\bЗАО': 'закрытое акционерное общество',
    r'\bОАО': 'открытое акционерное общество',
    r'\bИП': 'индивидуальный предприниматель',
    r'\bФГУП': 'федеральное государственное унитарное предприятие',
    r'\bГУП': 'государственное унитарное предприятие',
    r'\bМУП': 'муниципальное унитарное предприятие',
    
    # Государственные органы
    r'\bМВД': 'министерство внутренних дел',
    r'\bФСБ': 'федеральная служба безопасности',
    r'\bМЧС': 'министерство по чрезвычайным ситуациям',
    r'\bМинобр': 'министерство образования',
    r'\bМинздрав': 'министерство здравоохранения',
    r'\bМинфин': 'министерство финансов',
    r'\bМинтруд': 'министерство труда',
    r'\bМинэконом': 'министерство экономического развития',
    
    # Новостные сокращения
    r'\bСМИ': 'средства массовой информации',
    r'\bТВ': 'телевидение',
    r'\bРТР': 'российское телевидение и радио',
    r'\bИТАР': 'информационное телеграфное агентство россии',
    r'\bРИА': 'российское информационное агентство',
    r'\bТАСС': 'телеграфное агентство советского союза',
}

# Словарь для нормализации пунктуации
PUNCTUATION_MAP = {
    '…': '...',
    '–': '-',
    '—': '-',
    '«': '"',
    '»': '"',
    '„': '"',
    '“': '"',
    '”': '"',
    '"': '"',
    '‘': "'",
    '’': "'",
    '`': "'",
    '´': "'",
}


class UniversalPreprocessor:
    """Универсальный предпроцессор текста."""
    
    def __init__(self, config: Optional[PreprocessingConfig] = None):
        """
        Инициализация предпроцессора.
        
        Args:
            config: Конфигурация предобработки
        """
        self.config = config or PreprocessingConfig()
        self._compile_patterns()
    
    def _compile_patterns(self):
        """Компилирует регулярные выражения для ускорения работы."""
        self.patterns = {
            'url': RE_URL,
            'email': RE_EMAIL,
            'phone': RE_PHONE,
            'number': RE_NUM,
            'currency': RE_CURRENCY,
            'percent': RE_PERCENT,
            'date': RE_DATE,
        }
    
    def replace_special_tokens(self, text: str) -> str:
        """Заменяет специальные элементы на унифицированные токены."""
        if not text:
            return ""
        
        if self.config.replace_urls:
            text = self.patterns['url'].sub('<URL>', text)
        
        if self.config.replace_emails:
            text = self.patterns['email'].sub('<EMAIL>', text)
        
        if self.config.replace_numbers:
            text = self.patterns['phone'].sub('<PHONE>', text)
            text = self.patterns['currency'].sub('<CURRENCY>', text)
            text = self.patterns['percent'].sub('<PERCENT>', text)
            text = self.patterns['date'].sub('<DATE>', text)
            text = self.patterns['number'].sub('<NUM>', text)
        
        return text
    
    def expand_abbreviations(self, text: str) -> str:
        """Раскрывает сокращения."""
        if not self.config.expand_abbreviations or not text:
            return text
        
        for pattern, replacement in COMMON_ABBREVIATIONS.items():
            text = re.sub(pattern, replacement, text, flags=re.I)
        
        return text
    
    def normalize_punctuation(self, text: str) -> str:
        """Нормализует пунктуацию."""
        if not text:
            return ""
        
        if self.config.normalize_quotes:
            for old, new in PUNCTUATION_MAP.items():
                text = text.replace(old, new)
        
        if self.config.normalize_dashes:
            text = re.sub(r'[–—]', '-', text)
        
        if self.config.normalize_punctuation:
            # Нормализуем множественные точки
            text = re.sub(r'\.{3,}', '...', text)
            # Нормализуем множественные восклицательные знаки
            text = re.sub(r'!{2,}', '!!', text)
            # Нормализуем множественные вопросительные знаки
            text = re.sub(r'\?{2,}', '??', text)
        
        return text
    
    def normalize_spaces(self, text: str) -> str:
        """Нормализует пробелы."""
        if not self.config.normalize_spaces or not text:
            return text
        
        # Убираем лишние пробелы
        text = re.sub(r'\s+', ' ', text)
        # Убираем пробелы перед пунктуацией
        text = re.sub(r'\s+([.,;:!?])', r'\1', text)
        # Добавляем пробел после пунктуации, если его нет
        text = re.sub(r'([.,;:!?])([^\s])', r'\1 \2', text)
        
        return text.strip()
    
    def preprocess(self, text: str) -> str:
        """
        Выполняет полную предобработку текста.
        
        Args:
            text: Исходный текст
        
        Returns:
            Предобработанный текст
        """
        if not text:
            return ""
        
        # Заменяем специальные токены
        text = self.replace_special_tokens(text)
        
        # Раскрываем сокращения
        text = self.expand_abbreviations(text)
        
        # Нормализуем пунктуацию
        text = self.normalize_punctuation(text)
        
        # Нормализуем пробелы
        text = self.normalize_spaces(text)
        
        return text
    
    def preprocess_corpus(self, input_path: str, output_path: str) -> int:
        """
        Предобрабатывает корпус в формате JSONL.
        
        Args:
            input_path: Путь к исходному файлу
            output_path: Путь к выходному файлу
        
        Returns:
            Количество обработанных статей
        """
        import json
        
        processed_count = 0
        
        with open(input_path, 'r', encoding='utf-8') as infile, \
             open(output_path, 'w', encoding='utf-8') as outfile:
            
            for line in infile:
                line = line.strip()
                if not line:
                    continue
                
                try:
                    article = json.loads(line)
                    
                    # Предобрабатываем текст статьи
                    if 'text' in article:
                        article['text'] = self.preprocess(article['text'])
                    
                    # Предобрабатываем заголовок
                    if 'title' in article:
                        article['title'] = self.preprocess(article['title'])
                    
                    # Записываем предобработанную статью
                    outfile.write(json.dumps(article, ensure_ascii=False) + '\n')
                    processed_count += 1
                    
                except json.JSONDecodeError:
                    continue
        
        return processed_count


def create_preprocessing_pipeline(config: Optional[PreprocessingConfig] = None) -> UniversalPreprocessor:
    """
    Создает конвейер предобработки с заданной конфигурацией.
    
    Args:
        config: Конфигурация предобработки
    
    Returns:
        Настроенный предпроцессор
    """
    return UniversalPreprocessor(config)


if __name__ == "__main__":
    # Пример использования
    test_text = """
    Компания ООО "Тест" (ул. Ленина, д. 1) сообщила о результатах за 2023 г.
    Контакты: info@test.ru, +7(495)123-45-67, сайт www.test.com
    Цена: 1000 руб., рост на 15% по сравнению с прошлым годом.
    Дата: 15.03.2024, т.е. вчера.
    """
    
    # Создаем предпроцессор с настройками по умолчанию
    preprocessor = UniversalPreprocessor()
    
    # Предобрабатываем текст
    processed = preprocessor.preprocess(test_text)
    print("Предобработанный текст:")
    print(processed)
    
    # Пример с кастомной конфигурацией
    custom_config = PreprocessingConfig(
        replace_urls=True,
        replace_emails=True,
        replace_numbers=False,  # Не заменяем числа
        expand_abbreviations=True,
        normalize_punctuation=True
    )
    
    custom_preprocessor = UniversalPreprocessor(custom_config)
    custom_processed = custom_preprocessor.preprocess(test_text)
    print("\nС кастомной конфигурацией:")
    print(custom_processed)