#!/usr/bin/env python3
"""
Утилита для автоматического извлечения потенциальных тем и ключевых слов
из собранных сырых данных Telegram
"""

import os
import sys
import json
import glob
import re
from collections import Counter
import nltk
from pathlib import Path
import logging

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

logger = logging.getLogger(__name__)

# Директории
INPUT_DIR = os.getenv('OUTPUT_DIR', 'data')
RAW_DIR = f"{INPUT_DIR}/raw"
STOPWORDS_FILE = "stopwords.txt"

def download_nltk_data():
    """Скачивает необходимые данные для NLTK"""
    try:
        nltk.data.find('tokenizers/punkt')
    except LookupError:
        nltk.download('punkt')
    
    try:
        nltk.data.find('corpora/stopwords')
    except LookupError:
        nltk.download('stopwords')

def load_stopwords():
    """Загружает стоп-слова из файла или использует встроенные"""
    try:
        if os.path.exists(STOPWORDS_FILE):
            with open(STOPWORDS_FILE, 'r', encoding='utf-8') as f:
                return set(line.strip().lower() for line in f if line.strip())
        else:
            # Если файл не существует, используем стандартные стоп-слова NLTK
            from nltk.corpus import stopwords
            return set(stopwords.words('russian') + stopwords.words('english'))
    except Exception as e:
        logger.error(f"Ошибка при загрузке стоп-слов: {e}")
        return set()

def preprocess_text(text):
    """Предобработка текста: удаление спецсимволов, приведение к нижнему регистру"""
    # Удаляем ссылки
    text = re.sub(r'https?://\S+', '', text)
    # Удаляем упоминания
    text = re.sub(r'@\w+', '', text)
    # Удаляем хештеги
    text = re.sub(r'#\w+', '', text)
    # Удаляем спецсимволы и цифры
    text = re.sub(r'[^\w\s]', ' ', text)
    text = re.sub(r'\d+', ' ', text)
    # Приводим к нижнему регистру
    text = text.lower()
    # Удаляем лишние пробелы
    text = re.sub(r'\s+', ' ', text).strip()
    
    return text

def extract_keywords(texts, stopwords, min_length=4, max_length=20, top_n=30):
    """Извлекает ключевые слова из текстов"""
    # Объединяем все тексты в один
    all_text = ' '.join(texts)
    
    # Токенизируем
    tokens = nltk.word_tokenize(all_text)
    
    # Фильтруем стоп-слова и короткие слова
    filtered_tokens = [token for token in tokens 
                      if token.lower() not in stopwords 
                      and min_length <= len(token) <= max_length]
    
    # Считаем частоту слов
    word_counts = Counter(filtered_tokens)
    
    # Возвращаем top_n самых частых слов
    return [word for word, _ in word_counts.most_common(top_n)]

def extract_bigrams(texts, stopwords, min_length=4, top_n=20):
    """Извлекает биграммы из текстов"""
    all_tokens = []
    
    for text in texts:
        tokens = nltk.word_tokenize(text)
        # Фильтруем стоп-слова и короткие слова
        filtered_tokens = [token.lower() for token in tokens 
                          if token.lower() not in stopwords 
                          and len(token) >= min_length]
        all_tokens.extend(filtered_tokens)
    
    # Создаем биграммы
    bigrams = list(nltk.bigrams(all_tokens))
    
    # Считаем частоту биграмм
    bigram_counts = Counter(bigrams)
    
    # Возвращаем top_n самых частых биграмм в виде строк
    return [f"{w1} {w2}" for (w1, w2), _ in bigram_counts.most_common(top_n)]

def extract_topics_from_raw_data(min_messages=10):
    """Извлекает потенциальные темы из собранных сырых данных"""
    # Ищем все .json файлы в директории с сырыми данными
    raw_files = glob.glob(f"{RAW_DIR}/*.json")
    
    if not raw_files:
        logger.warning("Нет файлов для анализа")
        return None
    
    # Загружаем стоп-слова
    stopwords = load_stopwords()
    
    # Собираем все тексты
    all_texts = []
    
    for file in raw_files:
        try:
            with open(file, 'r', encoding='utf-8') as f:
                data = json.load(f)
            
            for msg in data["messages"]:
                if msg["text"]:
                    preprocessed = preprocess_text(msg["text"])
                    if preprocessed:  # Проверяем, что текст не пустой после предобработки
                        all_texts.append(preprocessed)
        except Exception as e:
            logger.error(f"Ошибка при чтении файла {file}: {e}")
    
    if len(all_texts) < min_messages:
        logger.warning(f"Недостаточно сообщений для анализа: {len(all_texts)} < {min_messages}")
        return None
    
    # Извлекаем ключевые слова
    keywords = extract_keywords(all_texts, stopwords)
    
    # Извлекаем биграммы
    bigrams = extract_bigrams(all_texts, stopwords)
    
    # Формируем предложения тем на основе биграмм и популярных слов
    suggested_topics = []
    
    for bigram in bigrams[:10]:  # Берем 10 лучших биграмм
        if bigram not in suggested_topics:
            suggested_topics.append(bigram)
    
    # Добавляем популярные ключевые слова, которые еще не в списке
    for keyword in keywords[:15]:  # Берем 15 лучших ключевых слов
        if keyword not in suggested_topics:
            suggested_topics.append(keyword)
    
    # Группируем слова по общей основе (очень упрощенный подход)
    grouped_topics = {}
    
    for topic in suggested_topics:
        # Используем первые 5 символов как приблизительную основу слова
        if len(topic) >= 5:
            root = topic[:5].lower()
            if root not in grouped_topics:
                grouped_topics[root] = []
            grouped_topics[root].append(topic)
    
    # Формируем финальный список тем и ключевых слов
    final_topics = {}
    
    for root, words in grouped_topics.items():
        # Выбираем самое длинное слово как название темы
        topic_name = max(words, key=len)
        
        # Остальные слова используем как ключевые слова
        topic_keywords = [word for word in words if word != topic_name]
        
        # Добавляем дополнительные ключевые слова из общего списка
        for keyword in keywords:
            if keyword.startswith(root[:3]) and keyword not in topic_keywords and keyword != topic_name:
                topic_keywords.append(keyword)
        
        final_topics[topic_name] = topic_keywords[:10]  # Ограничиваем 10 ключевыми словами
    
    return final_topics

def save_suggested_topics(topics):
    """Сохраняет предложенные темы в файл"""
    if not topics:
        logger.warning("Нет тем для сохранения")
        return False
    
    try:
        suggested_file = "suggested_topics.json"
        
        with open(suggested_file, 'w', encoding='utf-8') as f:
            json.dump(topics, f, indent=2, ensure_ascii=False)
        
        logger.info(f"Предложенные темы сохранены в {suggested_file}")
        
        # Также создаем более читаемую версию в Markdown
        md_file = "suggested_topics.md"
        
        with open(md_file, 'w', encoding='utf-8') as f:
            f.write("# Предложенные темы и ключевые слова\n\n")
            
            for topic, keywords in topics.items():
                f.write(f"## {topic.capitalize()}\n\n")
                f.write("Ключевые слова:\n")
                for keyword in keywords:
                    f.write(f"- {keyword}\n")
                f.write("\n")
        
        return True
    except Exception as e:
        logger.error(f"Ошибка при сохранении предложенных тем: {e}")
        return False

def main():
    """Основная функция для запуска экстрактора тем"""
    try:
        download_nltk_data()
        
        logger.info("Извлечение тем из собранных данных...")
        topics = extract_topics_from_raw_data()
        
        if topics:
            save_suggested_topics(topics)
            
            # Выводим результаты в консоль
            print("\nПредложенные темы и ключевые слова:")
            for topic, keywords in topics.items():
                print(f"\n{topic.capitalize()}:")
                print("Ключевые слова:")
                for keyword in keywords:
                    print(f"- {keyword}")
        else:
            print("Не удалось извлечь темы из данных. Возможно, недостаточно сообщений.")
    except Exception as e:
        logger.error(f"Ошибка: {e}")
        sys.exit(1)

if __name__ == "__main__":
    main()