"""Conservative repair of LaTeX escaping in raw model-generated JSON.

Call before JSON decoding, including for valid JSON: ``\\frac`` and ``\\not``
can otherwise be silently decoded as control characters. This is heuristic;
unknown commands starting with a valid JSON escape are left alone.
"""
from __future__ import annotations

import re

# Only these content fields carry TeX. Never rewrite labels or structural keys.
_CONTENT_FIELDS = {'question_content', 'solution_content', 'context_content', 'rubric_content'}
_CONTENT_KEYS = {f'"{field}"' for field in _CONTENT_FIELDS}
_STRING = re.compile(r'"(?:\\.|[^"\\])*"', re.DOTALL)
_COMMAND = re.compile(r'[a-zA-Z]+')
_MATH = re.compile(r"(\$\$.*?\$\$|\$.*?\$|\\\(.*?\\\)|\\\[.*?\\\])", re.DOTALL)
_DECODED_COMMAND = re.compile(r"[\b\f\n\r\t]([A-Za-z]+)")
_JSON_CONTROL_PREFIX = {"\b": "b", "\f": "f", "\n": "n", "\r": "r", "\t": "t"}


def normalize_literal_newlines(text: str) -> str:
    r"""Decode stray \n tokens in extracted TeX, preserving commands and \\.

    An odd backslash run ends in an escaped newline; any preceding pairs
    remain LaTeX row breaks. Even runs are already legitimate LaTeX.
    """
    def replace(match: re.Match[str]) -> str:
        slashes = match.group(1)
        return slashes[:-1] + "\n" if len(slashes) % 2 else match.group(0)

    return re.sub(r"(\\+)n(?![A-Za-z])", replace, text)
# Commands whose first letter is also a valid JSON escape. Other commands do
# not need a dictionary: their single backslash is invalid JSON already.
_AMBIGUOUS_COMMANDS = set('''
begin beta bar backslash binom boxed boldsymbol big bigl bigr biggl biggr
bigskip blacktriangle blacksquare bmod breve bullet bra braket
frac fbox forall footnote footnotesize framebox
not neq neg nabla newline newcommand newenvironment noindent nonumber
nexists notin nmid nparallel nsubseteq nsupseteq normalfont normalsize
rightarrow right leftarrow rangle rbrace rceil rfloor rightarrowtail
Rightarrow rightharpoonup rightharpoondown rightleftharpoons rightleftarrows
rm rule ref rho rightskip raisebox require renewcommand
text textbf textit textrm textsf texttt textnormal textcolor textsc textsl
textup textwidth theta tau tan tanh times to top triangle triangleleft
triangleright tilde tfrac today tag tiny thickspace thinspace tbinom
underbrace underline underset uparrow updownarrow upsilon usepackage
'''.split())


def repair_decoded_latex_controls(text: str) -> str:
    r"""Restore known TeX commands after JSON decoded their first letter.

    Only math spans and whitelisted command names are changed. This avoids
    treating ordinary paragraph newlines or tabs as damaged LaTeX.
    """
    def repair_math(math_match: re.Match[str]) -> str:
        def repair_command(command_match: re.Match[str]) -> str:
            candidate = (
                _JSON_CONTROL_PREFIX[command_match.group(0)[0]]
                + command_match.group(1)
            )
            return (
                f"\\{candidate}"
                if candidate in _AMBIGUOUS_COMMANDS
                else command_match.group(0)
            )

        return _DECODED_COMMAND.sub(repair_command, math_match.group(0))

    return _MATH.sub(repair_math, text)


def _repair_string(token: str) -> str:
    output = []
    index = 0
    while index < len(token):
        char = token[index]
        if char != '\\' or index + 1 == len(token):
            output.append(char)
            index += 1
            continue
        following = token[index + 1]
        # Consume already escaped pairs together, never double-escape them.
        if following in '\\"/':
            output.append(token[index:index + 2])
            index += 2
            continue
        command_match = _COMMAND.match(token, index + 1)
        command = command_match.group() if command_match else ''
        unicode_escape = following == 'u' and re.match(r'[0-9a-fA-F]{4}', token[index + 2:])
        if not unicode_escape and (following not in 'bfnrtu' or command in _AMBIGUOUS_COMMANDS
                                   or following == 'u'):
            output.append('\\\\')
        else:
            output.append('\\')
        index += 1
    return ''.join(output)


def repair_latex_json(
    raw: str, *, additional_content_fields: set[str] | None = None
) -> str:
    """Repair single TeX backslashes without decoding or rewriting other JSON.

    Preserve standard escapes, Unicode escapes, and already escaped TeX.
    A literal newline followed by e.g. 'ot' has the same wire representation
    as a broken ``\\not`` command; known TeX command names take precedence.
    Raw responses/caches should be retained to allow reviewing such cases.
    """
    additional_content_keys = {
        f'"{field}"' for field in (additional_content_fields or set())
    }
    output = []
    end = 0
    previous = None
    for match in _STRING.finditer(raw):
        token = match.group()
        if previous is not None and raw[previous.end():match.start()].strip() == ':':
            if previous.group() in _CONTENT_KEYS:
                token = _repair_string(token)
            elif previous.group() in additional_content_keys:
                token = _MATH.sub(
                    lambda math_match: _repair_string(math_match.group(0)),
                    token,
                )
        output.extend((raw[end:match.start()], token))
        end = match.end()
        previous = match
    output.append(raw[end:])
    return ''.join(output)