"""Server-side punctuation restoration for voice transcripts.

Uses simple rule-based approach (no external model dependency).
Can be replaced with ML-based model (e.g., silero, yandex punctuator) later.
"""

from __future__ import annotations

import re

# Words that typically start a new sentence
_SENTENCE_STARTERS: set[str] = {
    "а", "но", "и", "да", "вот", "так", "это", "тот", "кто", "что",
    "как", "когда", "где", "почему", "зачем", "сколько", "чей",
    "во-первых", "во-вторых", "наконец", "однако", "причем", "притом",
    "кстати", "например", "между", "тем", "впрочем", "значит",
    "итак", "следовательно", "кроме", "того", "помимо",
    "я", "ты", "он", "она", "оно", "мы", "вы", "они",
    "этот", "эта", "это", "эти", "мой", "твой", "наш", "ваш",
    "сегодня", "завтра", "вчера", "сейчас", "потом", "после",
    "сначала", "затем", "далее", "теперь", "тут", "там", "здесь",
}

# Interrogative words
_QUESTION_WORDS: set[str] = {
    "кто", "что", "какой", "какая", "какое", "какие", "чей", "чья", "чьё", "чьи",
    "сколько", "когда", "где", "куда", "откуда", "почему", "зачем", "как",
    "неужели", "разве", "ли",
}

# Exclamation words
_EXCLAMATION_WORDS: set[str] = {
    "ах", "ох", "ух", "ой", "эй", "ура", "браво", "караул",
    "здорово", "отлично", "прекрасно", "замечательно", "классно",
    "ужасно", "кошмар", "боже",
}


def restore_punctuation(text: str) -> str:
    """Add basic punctuation and capitalization to raw transcribed text.

    Works on Russian text from Whisper/STT output which typically
    comes without punctuation.
    """
    if not text or not text.strip():
        return text

    text = text.strip()

    # Split into clauses by common separators
    clauses = re.split(r"(?<=[^.!?])[;,]\s+", text)
    formatted: list[str] = []

    for clause in clauses:
        clause = clause.strip()
        if not clause:
            continue

        # Ensure first word is capitalized
        words = clause.split()
        if words:
            words[0] = words[0].capitalize()

        # Detect sentence type and add punctuation
        first_word = words[0].lower().rstrip(",.!?")

        # Remove any trailing punctuation
        last_word = words[-1].rstrip(",.!?") if words else ""

        if first_word in _EXCLAMATION_WORDS:
            punctuation = "!"
        elif first_word in _QUESTION_WORDS:
            punctuation = "?"
        elif last_word in _QUESTION_WORDS:
            punctuation = "?"
        else:
            punctuation = "."

        formatted.append(" ".join(words) + punctuation)

    result = " ".join(formatted)

    # Clean up common issues
    result = re.sub(r"\s+", " ", result)
    result = re.sub(r"\s*([.!?,;:])\s*", r"\1 ", result)
    result = re.sub(r"\s+\.", ".", result)
    result = result.strip()

    return result