from __future__ import annotations

import argparse
import re
from collections.abc import Sequence
from pathlib import Path
from typing import Any

from copienator import (
    EvaluationWorkspace,
    ExitCode,
    atomic_write_bytes,
    atomic_write_json,
    evaluation_parser,
    execute,
    read_json,
    workspace_from_args,
)

WORD_LIST_FILE = Path(__file__).resolve().parents[1] / "data" / "liste_francais.txt"
ACCENT_PATTERN = re.compile(r"[éèêëàâäîïôöùûüçœÉÈÊËÀÂÄÎÏÔÖÙÛÜÇŒ]")
MATH_PATTERN = re.compile(
    r"(\$\$.*?\$\$|\$.*?\$|\\\(.*?\\\)|\\\[.*?\\\])",
    re.DOTALL,
)


def build_parser() -> argparse.ArgumentParser:
    return evaluation_parser("Clean encoding and LaTeX issues in correction.json.")


def escape_latex_underscores(text: str) -> str:
    r"""Escape underscores outside math without double-escaping existing ones."""

    def escape_plain(value: str) -> str:
        # Collapse any existing escape run as well, making cleanup idempotent.
        return re.sub(r"\\*_", lambda _match: r"\_", value)

    parts: list[str] = []
    last_end = 0
    for match in MATH_PATTERN.finditer(text):
        start, end = match.span()
        parts.append(escape_plain(text[last_end:start]))
        parts.append(match.group(0))
        last_end = end
    parts.append(escape_plain(text[last_end:]))
    return "".join(parts)


def normalize_overescaped_latex_commands(text: str) -> str:
    r"""Collapse doubled command escapes inside LaTeX math environments.

    Model responses occasionally contain ``\\mathbb`` after JSON decoding where
    LaTeX requires ``\mathbb``.  A doubled backslash followed by whitespace is a
    legitimate row break (for example in ``cases``), so it must be preserved.
    """

    def normalize_math(match: re.Match[str]) -> str:
        return re.sub(r"\\\\(?=[A-Za-z{}])", r"\\", match.group(0))

    return MATH_PATTERN.sub(normalize_math, text)


def build_lookup_map(word_list_path: Path = WORD_LIST_FILE) -> dict[str, str]:
    words = word_list_path.read_text(encoding="utf-8").splitlines()
    lookup: dict[str, str] = {}
    for word in words:
        broken_key = ACCENT_PATTERN.sub("\x00", word)
        if "\x00" in broken_key:
            lookup[broken_key.lower()] = word
    return lookup


def fast_fix(text: str, lookup: dict[str, str]) -> str:
    def replacer(match: re.Match[str]) -> str:
        broken_word = match.group(0)
        return lookup.get(broken_word.lower(), broken_word)

    return re.sub(r"[a-zA-Z\x00]+", replacer, text)


def fix_hex_corruption_safe(text: str) -> str:
    return re.sub(
        r"\x00([eEfF][0-9a-fA-F])",
        lambda match: chr(int(match.group(1), 16)),
        text,
    )


def repair_json_escape_corruption(text: str) -> str:
    r"""Restore observed LaTeX commands consumed as JSON control escapes."""
    replacements = (
        ("\x0crac", r"\frac"),
        ("\x0ceuille", r"\equiv"),
        ("\theta", r"\theta"),
        ("\times", r"\times"),
        ("\textbackslash ", "\\"),
        ("\negthinspace", r"\negthinspace"),
        ("\neq", r"\neq"),
        ("\not", r"\not"),
        ("∈", r"\ensuremath{\in}"),
        ("⊂", r"\ensuremath{\subset}"),
    )
    for broken, repaired in replacements:
        text = text.replace(broken, repaired)
    return text


def clean_string(text: str, lookup: dict[str, str]) -> str:
    text = fix_hex_corruption_safe(text)
    text = text.replace("\x19", "\x00")
    text = text.replace("\x18", "\x00")
    text = text.replace("\x00\x00", "\x00")
    text = re.sub(r" \x00{1,2} ", " à ", text)
    if "\x00" in text:
        text = fast_fix(text, lookup).replace("\x00", "")
    text = repair_json_escape_corruption(text)
    text = normalize_overescaped_latex_commands(text)
    return escape_latex_underscores(text)


def clean_obj(value: Any, lookup: dict[str, str]) -> Any:
    if isinstance(value, str):
        return clean_string(value, lookup)
    if isinstance(value, list):
        return [clean_obj(item, lookup) for item in value]
    if isinstance(value, dict):
        return {
            key: item if key == "suffix" else clean_obj(item, lookup)
            for key, item in value.items()
        }
    return value


def run(
    workspace: EvaluationWorkspace,
    *,
    word_list_path: Path = WORD_LIST_FILE,
) -> ExitCode:
    workspace.require_files("correction.json")
    lookup = build_lookup_map(word_list_path)
    data = read_json(workspace.correction_file)
    cleaned = clean_obj(data, lookup)
    backup = workspace.root / "correction_precleanup.json"
    atomic_write_bytes(backup, workspace.correction_file.read_bytes())
    print(f"Original JSON backed up to {backup}")
    atomic_write_json(workspace.correction_file, cleaned)
    print(f"Fixed JSON saved to {workspace.correction_file}")
    return ExitCode.SUCCESS


def main(argv: Sequence[str] | None = None) -> int:
    parser = build_parser()
    return execute(parser, argv, lambda args: run(workspace_from_args(args)))


if __name__ == "__main__":
    raise SystemExit(main())