"""Утилиты обработки текста для конвейера."""
import logging
import re
from typing import Dict, List, Any

logger = logging.getLogger(__name__)


def filter_short_and_invalid_words(text: str) -> str:
    """Фильтрует текст, оставляя только валидные слова длиной более 2 символов
    на русском или английском языках.

    Используется для bm25 поиска
    """
    words = text.split()

    pattern = re.compile(r'^[а-яА-ЯёЁa-zA-Z]+')

    filtered_words = []
    for word in words:
        if len(word) > 1 and pattern.match(word) is not None:
            clean_word = re.sub(r'[\*\#\}\(\)\|\{\]\[\'\"\`\,\.\:\;\"\"\\]', '', word)
            filtered_words.append(clean_word)

    return ' '.join(filtered_words)

def get_few_shot_examples(
    prompt_data: Dict[str, Any], 
    example_project: str, 
    example_page: str, 
    is_few_shot: bool = True
) -> List[Dict[str, str]]:
    """Выбирает примеры для few-shot из данных промпта.
    
    Args:
        prompt_data: Данные промпта для категории
        example_project: Название проекта для примеров
        example_page: Название страницы для примеров
        is_few_shot: Флаг, включено ли few-shot обучение
        
    Returns:
        Список примеров для few-shot обучения
    """
    if not is_few_shot:
        return []
    
    return [
        ex
        for ex in prompt_data.get('examples', [])
        if ex.get('project') == example_project and ex.get('page') == example_page
    ]

def create_few_shot_prompt(
    system_message: str,
    examples: List[Dict[str, str]],
    user_question: str,
    template: str,
    context: str,
    previous_generations: str,
) -> List[Dict[str, str]]:
    """Формирует набор сообщений с few-shot в языковую модель.

    Args:
        system_message: Системное сообщение для LLM
        examples: Список примеров для few-shot обучения(формат  [{role: "user|assistant", content: "текст"}])
        user_question: Вопрос пользователя
        template: Шаблон для формирования ответа
        context: Контекстная информация
        previous_generations: Результаты предыдущих генераций

    Returns:
        messages: Список отформатированных сообщений для отправки в LLM формата
        [{role: "user|assistant|system", content: "текст"}]
    """
    messages = [{'role': 'system', 'content': system_message.strip()}]

    for example in examples:
        messages.append(
            {
                'role': 'user',
                'content': (
                    'ПРИМЕЧАНИЕ: Следующий текст - это только ПРИМЕР,'
                    'НЕ ИСПОЛЬЗУЙ конкретные данные из него для ответа на текущий запрос.\n'
                    f"Вопрос: {example['question']}"
                ).strip(),
            },
        )
        messages.append({'role': 'assistant', 'content': example['answer'].strip()})

    messages.append(
        {
            'role': 'user',
            'content': (
                f'Вопрос:\n{user_question}\n'
                f'Напиши по инструкции ниже:\n{template}\n'
                f'Основываясь на эти данные:\n{context.strip()}\n'
                f'Учитывая предыдущие генерации:\n{previous_generations}\n'
            ),
        },
    )

    return messages


def create_summarize_text_prompt(context_text: str) -> List[Dict[str, str]]:
    """Генерирует сообщение на саммари текста. С ограничением в 100 слов.

    Args:
        context_text: Текст для сокращения

    Returns:
        summary_messages: список сообщений в формате [{role: "user|system", content: "текст"}]
    """
    summary_messages = [
        {
            'role': 'system',
            'content': (
                'Ты должен сократить и саммаризировать текст, сохраняя только ключевую информацию. '
                'Обязательно сохрани все сущности'
                'Сокращённый текст должен быть не больше 100 слов.'
            ),
        },
        {'role': 'user', 'content': f'Сократи и саммаризируй следующий текст: \n{context_text}'},
    ]
    return summary_messages
