#!/usr/bin/env python3
"""
tts_preprocess.py — Intelligent text preprocessing for Kokoro TTS
# ADR: docs/adr/2026-02-19-tts-text-preprocessing.md

Solves the hard-wrap problem: text copied from terminals, AI chat output,
or PDF viewers has line breaks inserted at column width (~80-120 chars),
mid-sentence. These create unnatural pauses or cut-off speech in TTS.

Algorithm:
  1. Split into paragraph blocks on blank lines (preserved as real breaks)
  2. Within each block, heuristically join soft-wrapped lines
  3. Strip markdown noise (headers, bold, italic, code ticks)
  4. Replace symbols with TTS-friendly equivalents (∴ → "Therefore", etc.)

Soft-wrap detection heuristics (join if ANY match):
  - Next line starts with lowercase letter          → definite continuation
  - Next line starts with ","  ";"  "-"             → punctuation continuation
  - Current line is long (>60 chars) and next line is
    not a list item or heading                      → probable word-wrap

Keep break if ANY match:
  - Blank line (paragraph separator)
  - Next line starts with a list marker (-, *, •, 1., a.)
  - Next line has leading indentation (code/structure)
  - Current line ends with sentence punctuation (.!?) AND
    next line starts uppercase AND current line is short (<60 chars)
  - Both lines are very short (<35 chars)           → headers / labels

Usage:
  python tts_preprocess.py                  # reads stdin, writes stdout
  python tts_preprocess.py "text here"      # inline text arg
  echo "text" | python tts_preprocess.py   # pipe
"""

import re
import sys
import unicodedata


# ---------------------------------------------------------------------------
# Web copy-paste sanitisation (run BEFORE all other processing)
# ---------------------------------------------------------------------------
# Zero-width / invisible characters that survive HTML → plain text copy
_INVISIBLE_CHARS = re.compile(
    '['
    '\u00ad'          # Soft hyphen
    '\u200b'          # Zero-width space
    '\u200c'          # Zero-width non-joiner
    '\u200d'          # Zero-width joiner
    '\u200e\u200f'    # LTR / RTL marks
    '\u2028\u2029'    # Line / paragraph separators
    '\u202a-\u202e'   # Bidi embedding controls
    '\u2060'          # Word joiner
    '\u2066-\u2069'   # Bidi isolate controls
    '\ufeff'          # BOM / zero-width no-break space
    '\ufffc'          # Object replacement character (inline images)
    ']'
)

# Whitespace variants → normal space
_FANCY_SPACES = re.compile(
    '['
    '\u00a0'          # Non-breaking space
    '\u2002'          # En space
    '\u2003'          # Em space
    '\u2007'          # Figure space
    '\u2009'          # Thin space
    '\u200a'          # Hair space
    '\u202f'          # Narrow no-break space
    '\u205f'          # Medium mathematical space
    ']'
)

# Smart quotes & typographic ligatures → ASCII equivalents
_SMART_QUOTES: list[tuple[str, str]] = [
    ('\u201c', '"'), ('\u201d', '"'),   # " "  → "
    ('\u2018', "'"), ('\u2019', "'"),   # ' '  → '
    ('\u2013', '-'),                     # –    → -  (en-dash)
    ('\ufb01', 'fi'), ('\ufb02', 'fl'), # fi/fl ligatures
    ('\ufb03', 'ffi'), ('\ufb04', 'ffl'),
]


def _sanitize_clipboard(text: str) -> str:
    """Strip invisible chars, normalise whitespace variants, fix smart quotes.

    Must run BEFORE markdown/symbol substitutions so patterns match clean text.
    """
    # Normalise Unicode (NFC) — decomposed chars → composed
    text = unicodedata.normalize('NFC', text)
    # Strip invisible / zero-width chars
    text = _INVISIBLE_CHARS.sub('', text)
    # Normalise fancy spaces → regular space
    text = _FANCY_SPACES.sub(' ', text)
    # Smart quotes / ligatures → ASCII
    for fancy, plain in _SMART_QUOTES:
        text = text.replace(fancy, plain)
    # Normalise line endings (CRLF / CR → LF)
    text = text.replace('\r\n', '\n').replace('\r', '\n')
    # Strip control chars except \n and \t
    text = re.sub(r'[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]', '', text)
    # Strip surrogate codepoints (cause 'surrogates not allowed' in phonemizer)
    text = re.sub('[\ud800-\udfff]', '', text)
    # Strip excessive combining marks / Zalgo text — keep at most 3
    text = re.sub('([\u0300-\u036f\u0489\u20d0-\u20ff\ufe20-\ufe2f]{3})[\u0300-\u036f\u0489\u20d0-\u20ff\ufe20-\ufe2f]+', r'\1', text)
    return text


# ---------------------------------------------------------------------------
# Symbol / markdown substitution table (order matters — regex applied top-down)
# ---------------------------------------------------------------------------
_SUBSTITUTIONS: list[tuple[str | re.Pattern, str]] = [
    # Markdown structural noise
    (re.compile(r'^∴\s*Thinking…?\s*$', re.MULTILINE), 'Thinking.'),  # "∴ Thinking…" → spoken
    (re.compile(r'^#{1,6}\s+', re.MULTILINE), ''),        # ATX headings
    (re.compile(r'^[-*_]{3,}\s*$', re.MULTILINE), ''),    # horizontal rules
    (re.compile(r'^>\s+', re.MULTILINE), ''),             # blockquotes
    # Inline formatting (must come before raw-symbol replacements)
    (re.compile(r'\*\*([^*\n]+)\*\*'), r'\1'),            # **bold**
    (re.compile(r'\*([^*\n]+)\*'), r'\1'),                # *italic*
    (re.compile(r'__([^_\n]+)__'), r'\1'),                # __bold__
    (re.compile(r'_([^_\n]+)_'), r'\1'),                  # _italic_
    (re.compile(r'`{3}[^\n]*\n', re.MULTILINE), ''),      # code fence open
    (re.compile(r'`{3}', re.MULTILINE), ''),              # code fence close
    (re.compile(r'`([^`\n]+)`'), r'\1'),                  # `inline code`
    # Markdown links: [text](url) → text
    (re.compile(r'\[([^\]]+)\]\([^)]+\)'), r'\1'),
    # Unicode symbols → TTS-readable words
    ('∴', 'Therefore,'),
    ('∵', 'Because,'),
    ('⇒', 'implies'),
    ('→', 'leads to'),
    ('←', 'comes from'),
    ('↔', 'is equivalent to'),
    ('≡', 'is equivalent to'),
    ('≈', 'approximately'),
    ('≠', 'is not equal to'),
    ('≤', 'is at most'),
    ('≥', 'is at least'),
    ('∈', 'in'),
    ('∉', 'not in'),
    ('∩', 'intersect'),
    ('∪', 'union'),
    ('∞', 'infinity'),
    ('✓', 'yes'),
    ('✗', 'no'),
    ('✘', 'no'),
    ('•', '-'),
    ('…', '...'),
    # Numbered list noise: "  4. Item" → "Item" (strip leading number+dot)
    # We keep the text; Kokoro reads "4." as "four period" otherwise
    (re.compile(r'^\s{0,4}\d{1,2}\.\s+', re.MULTILINE), ''),
    # Lettered list: "  a. Item" → "Item"
    (re.compile(r'^\s{0,4}[a-z]\.\s+', re.MULTILINE), ''),
    # Dash/bullet list markers at line start — keep a space for prosody
    (re.compile(r'^\s*[-*•]\s+', re.MULTILINE), ' '),
    # Box chars from tables/diagrams → skip
    (re.compile(r'[│├└─┌┐┘╔╗╚╝║═╠╣╦╩╬]'), ' '),
    # Multiple spaces → single (only mid-line, preserve leading indentation for _INDENTED check)
    (re.compile(r'(?<=\S)[ \t]{2,}'), ' '),
]

# Sentence-ending punctuation (for keep-break heuristic)
# List-item start: -, *, •, 1., a. (must be at start of stripped line)
_LIST_START = re.compile(r'^(\s{0,4}[-*•]|\s{0,4}\d{1,2}\.\s|\s{0,4}[a-z]\.\s)')
# Leading indentation (code blocks, deeply nested structure)
_INDENTED = re.compile(r'^\s{4,}')
# Very short line threshold (probable header or label)
_SHORT_LINE = 35
# Long line threshold — heuristic for soft-wrap detection
_LONG_LINE = 62


def _apply_substitutions(text: str) -> str:
    for pattern, replacement in _SUBSTITUTIONS:
        if isinstance(pattern, str):
            text = text.replace(pattern, replacement)
        else:
            text = pattern.sub(replacement, text)
    return text


def _should_join(prev_original_len: int, current: str, nxt: str) -> bool:
    """Return True if `nxt` should be joined onto `current` with a space.

    Args:
        prev_original_len: length of the last original line that was appended
            (NOT the accumulated result length — used for soft-wrap detection).
        current: accumulated result so far.
        nxt: the next line to consider joining.
    """
    c = current.rstrip()
    n = nxt.strip()

    if not n:
        return False  # blank line → paragraph break (caller handles)

    # Preserve list items, indented blocks, headings
    if _LIST_START.match(n):
        return False
    if _INDENTED.match(nxt):
        return False
    if n.startswith('#'):
        return False

    # Current line ends with sentence-ending punctuation (.!?) AND next starts
    # uppercase → real paragraph/sentence break.
    # Exclude colon — "Note: The system" is not a sentence break.
    if re.search(r'[.!?]\s*$', c) and n and n[0].isupper():
        return False

    # Clear continuation: next starts with lowercase
    if n and n[0].islower():
        return True

    # Continuation punctuation at start of next line
    if n and n[0] in (',', ';', ')'):
        return True

    # Both lines are very short → intentional structure (e.g., step labels)
    if prev_original_len < _SHORT_LINE and len(n) < _SHORT_LINE:
        return False

    # Soft-wrap detection: if the previous original line was near terminal
    # column width (62-200 chars), it was likely soft-wrapped
    if _LONG_LINE <= prev_original_len <= 200:
        return True

    return False


def _reflow_block(lines: list[str]) -> str:
    """Join soft-wrapped lines within a single paragraph block."""
    if not lines:
        return ''

    result = lines[0].rstrip()
    last_orig_len = len(result)
    for i in range(1, len(lines)):
        nxt = lines[i]
        if _should_join(last_orig_len, result, nxt):
            # Avoid double-space if result already ends with space
            result = result.rstrip() + ' ' + nxt.strip()
            # Don't update last_orig_len — the joined line extends the same
            # original line, so soft-wrap detection should use the trigger line
        else:
            result = result + '\n' + nxt.rstrip()
            last_orig_len = len(nxt.rstrip())
    return result


def _split_sentences(text: str) -> str:
    """
    Optional: split reflowed paragraphs into one-sentence-per-line using pySBD.
    pySBD handles abbreviations (Dr., p. 5, etc.) that fool naive regex.
    Falls back gracefully if pySBD is not installed.
    """
    try:
        import pysbd  # type: ignore[import]
        seg = pysbd.Segmenter(language='en', clean=False)
        paragraphs = text.split('\n\n')
        result_paragraphs = []
        for para in paragraphs:
            if para.strip():
                sentences = seg.segment(para)
                result_paragraphs.append(' '.join(s.strip() for s in sentences))
            else:
                result_paragraphs.append('')
        return '\n\n'.join(result_paragraphs)
    except ImportError:
        return text


def preprocess(text: str, split_sentences: bool = False) -> str:
    """Full TTS preprocessing pipeline."""
    # 0. Sanitise web copy-paste artifacts (invisible chars, smart quotes, etc.)
    text = _sanitize_clipboard(text)

    # 1. Apply symbol / markdown substitutions
    text = _apply_substitutions(text)

    # 2. Split into paragraph blocks (one or more blank lines = separator)
    #    We preserve blank lines as paragraph separators in output.
    raw_blocks = re.split(r'\n\s*\n', text)

    processed_blocks: list[str] = []
    for block in raw_blocks:
        lines = block.split('\n')
        reflowed = _reflow_block([ln for ln in lines])
        if reflowed.strip():
            processed_blocks.append(reflowed.strip())

    # 3. Rejoin with double-newline paragraph separators
    result = '\n\n'.join(processed_blocks)

    # 4. Final cleanup: collapse 3+ newlines, trailing whitespace
    result = re.sub(r'\n{3,}', '\n\n', result)
    result = re.sub(r'[ \t]+\n', '\n', result)

    # 5. Optional: sentence splitting (pySBD) for more natural TTS pacing
    if split_sentences:
        result = _split_sentences(result)

    return result.strip()


def main() -> None:
    split = '--split-sentences' in sys.argv
    args = [a for a in sys.argv[1:] if not a.startswith('--')]

    if args:
        # Inline text passed as argument
        text = ' '.join(args)
    elif not sys.stdin.isatty():
        text = sys.stdin.read()
    else:
        print(__doc__)
        sys.exit(0)

    print(preprocess(text, split_sentences=split), end='')


if __name__ == '__main__':
    main()
