# src/text_cleaner.py
"""
Модуль для очистки и предобработки текста.
Выполняет удаление HTML-разметки, служебных символов, рекламных блоков,
стандартизацию пробельных символов и фильтрацию стоп-слов.
"""

import re
from typing import List, Optional
from bs4 import BeautifulSoup
import nltk
from nltk.corpus import stopwords

# Загружаем русские стоп-слова
try:
    RU_STOP = set(stopwords.words('russian'))
except LookupError:
    nltk.download('stopwords')
    RU_STOP = set(stopwords.words('russian'))

# Дополнительные стоп-слова для новостных текстов
NEWS_STOP_WORDS = {
    'сообщает', 'сообщил', 'сообщила', 'сообщили', 'сообщило',
    'заявил', 'заявила', 'заявили', 'заявило',
    'отметил', 'отметила', 'отметили', 'отметило',
    'подчеркнул', 'подчеркнула', 'подчеркнули', 'подчеркнуло',
    'уточнил', 'уточнила', 'уточнили', 'уточнило',
    'добавил', 'добавила', 'добавили', 'добавило',
    'пояснил', 'пояснила', 'пояснили', 'пояснило',
    'сказал', 'сказала', 'сказали', 'сказало',
    'говорит', 'говорят', 'говорил', 'говорила',
    'пишет', 'пишут', 'писал', 'писала',
    'читайте', 'также', 'также', 'также',
    'подробнее', 'далее', 'продолжение', 'следует'
}

RU_STOP.update(NEWS_STOP_WORDS)


def remove_html(text: str) -> str:
    """Удаляет HTML-разметку из текста."""
    if not text:
        return ""
    soup = BeautifulSoup(text, 'html.parser')
    return soup.get_text(separator=' ')


def normalize_whitespace(text: str) -> str:
    """Стандартизирует пробельные символы."""
    if not text:
        return ""
    # Заменяем все виды пробелов на обычные
    text = re.sub(r'[\s\u00A0\u2000-\u200B\u2028\u2029\u202F\u205F\u3000]+', ' ', text)
    return text.strip()


def remove_nontext_chars(text: str) -> str:
    """Удаляет служебные символы, оставляя кириллицу, латиницу и пунктуацию."""
    if not text:
        return ""
    # Оставляем буквы, цифры, пробелы и основную пунктуацию
    return re.sub(r'[^\w\s\-\.,;:\?!\'"«»()—–№]', ' ', text)


def remove_stopwords_tokens(tokens: List[str]) -> List[str]:
    """Удаляет стоп-слова из списка токенов."""
    if not tokens:
        return []
    return [t for t in tokens if t.lower() not in RU_STOP and len(t.strip()) > 0]


def remove_short_tokens(tokens: List[str], min_length: int = 2) -> List[str]:
    """Удаляет слишком короткие токены."""
    if not tokens:
        return []
    return [t for t in tokens if len(t.strip()) >= min_length]


def remove_numeric_tokens(tokens: List[str]) -> List[str]:
    """Удаляет токены, состоящие только из цифр."""
    if not tokens:
        return []
    return [t for t in tokens if not t.isdigit()]


def clean_text(text: str, 
               lower: bool = True, 
               remove_stopwords: bool = False,
               min_token_length: int = 2,
               remove_numbers: bool = False) -> str:
    """
    Основная функция очистки текста.
    
    Args:
        text: Исходный текст
        lower: Приводить к нижнему регистру
        remove_stopwords: Удалять стоп-слова
        min_token_length: Минимальная длина токена
        remove_numbers: Удалять числовые токены
    
    Returns:
        Очищенный текст
    """
    if not text:
        return ""
    
    # Удаляем HTML
    text = remove_html(text)
    
    # Нормализуем пробелы
    text = normalize_whitespace(text)
    
    # Приводим к нижнему регистру
    if lower:
        text = text.lower()
    
    # Удаляем служебные символы
    text = remove_nontext_chars(text)
    
    # Нормализуем пробелы еще раз
    text = normalize_whitespace(text)
    
    # Если нужно удалить стоп-слова или числа, токенизируем
    if remove_stopwords or remove_numbers:
        tokens = text.split()
        
        if remove_stopwords:
            tokens = remove_stopwords_tokens(tokens)
        
        if remove_numbers:
            tokens = remove_numeric_tokens(tokens)
        
        if min_token_length > 1:
            tokens = remove_short_tokens(tokens, min_token_length)
        
        text = ' '.join(tokens)
    
    return text


def clean_corpus_jsonl(input_path: str, 
                      output_path: str,
                      **clean_kwargs) -> int:
    """
    Очищает корпус в формате JSONL.
    
    Args:
        input_path: Путь к исходному файлу
        output_path: Путь к выходному файлу
        **clean_kwargs: Параметры для clean_text
    
    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'] = clean_text(article['text'], **clean_kwargs)
                
                # Очищаем заголовок
                if 'title' in article:
                    article['title'] = clean_text(article['title'], **clean_kwargs)
                
                # Записываем очищенную статью
                outfile.write(json.dumps(article, ensure_ascii=False) + '\n')
                processed_count += 1
                
            except json.JSONDecodeError:
                continue
    
    return processed_count


if __name__ == "__main__":
    # Пример использования
    test_text = """
    <p>Это <strong>тестовый</strong> текст с HTML-разметкой.</p>
    <br/>Он содержит    множественные    пробелы   и
    различные символы: @#$%^&*().
    """
    
    cleaned = clean_text(test_text, lower=True, remove_stopwords=False)
    print("Очищенный текст:", cleaned)