"""ANSI escape code handling utilities.

Provides helpers to strip ANSI/VT escape sequences and control characters
so that prompt-pattern matching operates on the *visible* text only.
"""

import re

# CSI cursor-forward (ESC [ n C): replace with n spaces so that TUI tools
# that use cursor movement for word spacing (e.g. Claude Code ≥ 2.x) produce
# readable text after cleaning.  Must be applied *before* the general strip.
_CURSOR_FORWARD_RE = re.compile(r"\x1b\[(\d*)C")

_ANSI_RE = re.compile(
    r"\x1b"
    r"(?:"
    r"\[[0-9;?]*[A-Za-z~]"  # CSI sequences (e.g. colors, cursor movement)
    r"|\][^\x07]*(?:\x07|\x1b\\)"  # OSC sequences (e.g. window title)
    r"|[()#][A-Za-z0-9]"  # character set / font selection
    r"|[A-Za-z=<>]"  # simple two-char sequences
    r")"
)

_CONTROL_RE = re.compile(r"[\x00-\x08\x0b\x0c\x0e-\x1f]")


def _expand_cursor_forward(m):
    """Substitute ESC[nC with *n* spaces (default 1 when n is omitted)."""
    n = m.group(1)
    return " " * (int(n) if n else 1)


def strip_ansi(text):
    """Remove ANSI / VT escape sequences from *text*.

    Cursor-forward sequences (``ESC [ n C``) are expanded to the equivalent
    number of spaces *before* other sequences are stripped.  This preserves
    word spacing in TUI tools (such as Claude Code) that render menus by
    advancing the cursor rather than emitting literal space characters.
    """
    text = _CURSOR_FORWARD_RE.sub(_expand_cursor_forward, text)
    return _ANSI_RE.sub("", text)


def strip_control(text):
    """Remove control characters except ``\\n`` (0x0a) and ``\\r`` (0x0d)."""
    return _CONTROL_RE.sub("", text)


def clean_text(text):
    """Strip escapes, resolve carriage-returns and produce clean visible text.

    ``\\r\\n`` (standard PTY line ending) is normalised to ``\\n`` first.
    Remaining bare ``\\r`` is treated as a terminal overwrite: the last
    *non-empty* carriage-return segment on each line wins.  Using the last
    non-empty (rather than the absolute last) segment correctly handles TUI
    tools (such as Claude Code ≥ 2.x) that emit ``\\r\\r\\n`` as a line
    terminator — the trailing bare ``\\r`` would otherwise erase the line
    content by making the final segment empty.
    """
    text = strip_ansi(text)
    text = strip_control(text)
    text = text.replace("\r\n", "\n")

    cleaned_lines = []
    for line in text.split("\n"):
        segments = line.split("\r")
        # Last non-empty segment is the visible line content.
        # Fall back to "" only if every segment is empty (blank line).
        result = next((s for s in reversed(segments) if s), "")
        cleaned_lines.append(result)

    return "\n".join(cleaned_lines)
