#!/usr/bin/env python3
"""
insert_d_into_nb.py

Рекурсивно обходит Wolfram Notebook (*.nb), пропускает файлы с
"фитирование" в имени и меняет значение d в первой Input-ячейке
на значение из файла коэффициентов.

Ключевая мера безопасности:
- notebook читается и пишется как bytes, без перекодирования и нормализации EOL;
- Internal cache information не удаляется и не редактируется;
- размер изменяемой первой Input-ячейки и размер всего .nb сохраняются
  побайтно неизменными: если новое число длиннее/короче старого,
  разница компенсируется незначащими пробелами в этой же Input-ячейке;
- перед любой записью создаётся backup в уникальной папке /tmp;
- запись атомарная, после неё выполняется повторная проверка.
"""

from __future__ import annotations

import argparse
import os
import re
import shutil
import stat
import sys
import tempfile
from dataclasses import dataclass
from decimal import Decimal, InvalidOperation
from pathlib import Path


SECTION_RE = re.compile(r"^\s*(\d+-\d+_\d+-\d+)\s*$")
D_LINE_RE = re.compile(
    r"^\s*d\s*=\s*([+-]?(?:\d+(?:\.\d*)?|\.\d+))\s*$"
)

# Примеры:
# h_154-86_Pu240_8-9_gid_Qw1.nb
# h_152-88_h_Pu240_6-7_pr_Qw1.nb
FILENAME_KEY_RE = re.compile(
    r"(?P<ab>\d+-\d+)(?:_h)?_Pu240_(?P<ij>\d+-\d+)",
    re.IGNORECASE,
)

BEGIN_MARKER = b"(* Beginning of Notebook Content *)"
END_MARKER = b"(* End of Notebook Content *)"

# Ищем именно Cell[BoxData[...], "Input"] и не разрешаем матчу
# перескочить через начало следующей Cell[...].
FIRST_INPUT_RE = re.compile(
    rb'Cell\[\s*BoxData\[\s*'
    rb'(?P<box>(?:(?!\bCell\[).)*?)'
    rb'\]\s*,\s*"Input"',
    re.DOTALL,
)

# Для первой переменной принимается консервативная и ожидаемая по ТЗ форма:
# RowBox[{"d", "=", "2.88"}]
DIRECT_ASSIGNMENT_RE = re.compile(
    rb'^\s*RowBox\[\{\s*'
    rb'"(?P<var>[A-Za-z$][A-Za-z0-9$]*)"\s*,\s*'
    rb'"="\s*,\s*'
    rb'"(?P<value>[+-]?(?:\d+(?:\.\d*)?|\.\d+))"\s*'
    rb'\}\]\s*$',
    re.DOTALL,
)


@dataclass(frozen=True)
class FirstAssignment:
    variable: str
    value_text: str
    value_start: int
    value_end: int
    first_input_end: int
    next_cell_start: int


class NotebookFormatError(RuntimeError):
    pass


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=(
            "Заменяет d в первой Input-ячейке Wolfram Notebook "
            "по значениям из файла коэффициентов."
        )
    )
    parser.add_argument(
        "--coeff_file",
        required=True,
        type=Path,
        help="Путь к txt-файлу коэффициентов.",
    )
    parser.add_argument(
        "--wolfram_notebooks",
        "--wolfram_noteooks",  # алиас для опечатки, встречающейся в ТЗ
        dest="wolfram_notebooks",
        required=True,
        type=Path,
        help="Корневая директория, где рекурсивно искать *.nb.",
    )
    return parser.parse_args()


def parse_coefficients(path: Path) -> dict[str, str]:
    try:
        text = path.read_text(encoding="utf-8-sig")
    except OSError as exc:
        raise RuntimeError(f"Не удалось прочитать файл коэффициентов: {exc}") from exc

    result: dict[str, str] = {}
    current_section: str | None = None
    seen_sections: set[str] = set()

    for lineno, raw_line in enumerate(text.splitlines(), start=1):
        section_match = SECTION_RE.match(raw_line)
        if section_match:
            current_section = section_match.group(1)
            if current_section in seen_sections:
                raise RuntimeError(
                    f"{path}:{lineno}: повторная секция {current_section!r}"
                )
            seen_sections.add(current_section)
            continue

        d_match = D_LINE_RE.match(raw_line)
        if d_match and current_section is not None:
            if current_section in result:
                raise RuntimeError(
                    f"{path}:{lineno}: повторное d в секции {current_section!r}"
                )
            value = d_match.group(1)

            # Дополнительная валидация числа.
            try:
                Decimal(value)
            except InvalidOperation as exc:
                raise RuntimeError(
                    f"{path}:{lineno}: некорректное значение d={value!r}"
                ) from exc

            result[current_section] = value

    if not result:
        raise RuntimeError(
            f"В {path} не найдено ни одной секции с корректной строкой d = ..."
        )

    missing_d = sorted(seen_sections - result.keys())
    if missing_d:
        raise RuntimeError(
            "В следующих секциях отсутствует корректная строка d: "
            + ", ".join(missing_d)
        )

    return result


def coefficient_key_from_filename(name: str) -> str | None:
    match = FILENAME_KEY_RE.search(name)
    if not match:
        return None
    return f"{match.group('ab')}_{match.group('ij')}"


def find_first_assignment(data: bytes) -> FirstAssignment:
    begin = data.find(BEGIN_MARKER)
    end = data.find(END_MARKER)

    if begin < 0 or end < 0 or end <= begin:
        raise NotebookFormatError(
            "не найдены корректные маркеры начала/конца Notebook Content"
        )

    content_start = begin + len(BEGIN_MARKER)
    content = data[content_start:end]

    input_match = FIRST_INPUT_RE.search(content)
    if not input_match:
        raise NotebookFormatError("не найдена первая Input-ячейка")

    box = input_match.group("box")
    assignment = DIRECT_ASSIGNMENT_RE.fullmatch(box)
    if not assignment:
        raise NotebookFormatError(
            "первая Input-ячейка не является простым присваиванием вида "
            'RowBox[{"d", "=", "число"}]'
        )

    var = assignment.group("var").decode("ascii")
    value = assignment.group("value").decode("ascii")

    box_abs_start = content_start + input_match.start("box")
    value_start = box_abs_start + assignment.start("value")
    value_end = box_abs_start + assignment.end("value")
    first_input_end = content_start + input_match.end()

    # Компенсацию длины делаем до следующей Cell[, то есть внутри
    # опций первой Input-ячейки и до начала следующей ячейки.
    next_cell_start = data.find(b"Cell[", first_input_end, end)
    if next_cell_start < 0:
        raise NotebookFormatError(
            "после первой Input-ячейки не найдена следующая Cell"
        )

    return FirstAssignment(
        variable=var,
        value_text=value,
        value_start=value_start,
        value_end=value_end,
        first_input_end=first_input_end,
        next_cell_start=next_cell_start,
    )


def decimal_equal(left: str, right: str) -> bool:
    try:
        return Decimal(left) == Decimal(right)
    except InvalidOperation:
        return left == right


def _remove_indentation_spaces(data: bytes, count: int) -> bytes:
    """
    Удаляет ровно count пробелов/табов только из отступов в начале строк.

    Переводы строк не меняются, содержимое строковых литералов не трогается.
    Для типичного CellChangeTimes в первой ячейке таких пробелов значительно
    больше, чем может понадобиться для компенсации длины d.
    """
    if count <= 0:
        return data

    out = bytearray()
    i = 0
    remaining = count
    at_line_start = False

    while i < len(data):
        b = data[i]

        if at_line_start and b in (0x20, 0x09) and remaining > 0:
            # Удаляем отступ по одному байту.
            remaining -= 1
            i += 1
            continue

        out.append(b)

        if b == 0x0A:  # \n; CR из CRLF уже был скопирован
            at_line_start = True
        elif b not in (0x20, 0x09, 0x0D):
            at_line_start = False

        i += 1

    if remaining:
        raise NotebookFormatError(
            "не хватило безопасных отступов для сохранения размера notebook "
            f"(нужно было убрать ещё {remaining} байт)"
        )

    return bytes(out)


def patch_d_preserving_size(
    original: bytes,
    assignment: FirstAssignment,
    new_value: str,
) -> bytes:
    new_value_bytes = new_value.encode("ascii")
    old_value_bytes = original[assignment.value_start:assignment.value_end]

    if old_value_bytes.decode("ascii") != assignment.value_text:
        raise NotebookFormatError("внутренняя проверка позиции старого d не прошла")

    before = original[:assignment.value_start]
    after_value = original[assignment.value_end:assignment.next_cell_start]
    tail = original[assignment.next_cell_start:]

    delta = len(new_value_bytes) - len(old_value_bytes)

    if delta > 0:
        # Новое число длиннее. Убираем столько же незначащих пробелов
        # из отступов ПОСЛЕ значения d, но до следующей Cell.
        compensated_after = _remove_indentation_spaces(after_value, delta)
    elif delta < 0:
        # Новое число короче. Добавляем пробелы сразу ПОСЛЕ закрывающей
        # кавычки числового RowBox-токена. Это Wolfram-синтаксически
        # незначащий whitespace.
        need = -delta
        if not after_value.startswith(b'"'):
            raise NotebookFormatError(
                "неожиданная структура после числового токена d"
            )
        compensated_after = (
            after_value[:1] + (b" " * need) + after_value[1:]
        )
    else:
        compensated_after = after_value

    patched = before + new_value_bytes + compensated_after + tail

    if len(patched) != len(original):
        raise NotebookFormatError(
            "после компенсации изменился размер notebook — запись отменена"
        )

    # Всё после начала следующей Cell должно остаться байт-в-байт неизменным.
    if patched[assignment.next_cell_start:] != original[assignment.next_cell_start:]:
        raise NotebookFormatError(
            "обнаружено изменение за пределами первой Input-ячейки"
        )

    # Cache-блок должен остаться полностью неизменным.
    old_cache_pos = original.find(b"(* Internal cache information *)", end_marker_pos(original))
    new_cache_pos = patched.find(b"(* Internal cache information *)", end_marker_pos(patched))
    if old_cache_pos < 0 or new_cache_pos < 0:
        raise NotebookFormatError("не найден завершающий Internal cache information")
    if old_cache_pos != new_cache_pos:
        raise NotebookFormatError("сместилась позиция завершающего cache-блока")
    if patched[old_cache_pos:] != original[old_cache_pos:]:
        raise NotebookFormatError("cache-блок был изменён — запись отменена")

    # Повторно разбираем результат и убеждаемся, что d стал целевым.
    check = find_first_assignment(patched)
    if check.variable != "d" or check.value_text != new_value:
        raise NotebookFormatError(
            "контрольная проверка после замены d не прошла"
        )

    return patched


def end_marker_pos(data: bytes) -> int:
    pos = data.find(END_MARKER)
    return pos if pos >= 0 else 0


def backup_file(source: Path, root: Path, backup_root: Path) -> Path:
    relative = source.relative_to(root)
    destination = backup_root / relative
    destination.parent.mkdir(parents=True, exist_ok=True)
    shutil.copy2(source, destination)
    return destination


def atomic_write(path: Path, data: bytes) -> None:
    original_stat = path.stat()
    temp_name: str | None = None

    try:
        with tempfile.NamedTemporaryFile(
            mode="wb",
            prefix=f".{path.name}.",
            suffix=".tmp",
            dir=path.parent,
            delete=False,
        ) as tmp:
            temp_name = tmp.name
            tmp.write(data)
            tmp.flush()
            os.fsync(tmp.fileno())

        os.chmod(temp_name, stat.S_IMODE(original_stat.st_mode))
        os.replace(temp_name, path)
        temp_name = None

        # На Linux fsync директории закрепляет rename в файловой системе.
        try:
            dir_fd = os.open(path.parent, os.O_RDONLY | os.O_DIRECTORY)
        except OSError:
            dir_fd = None
        if dir_fd is not None:
            try:
                os.fsync(dir_fd)
            finally:
                os.close(dir_fd)

    finally:
        if temp_name is not None:
            try:
                os.unlink(temp_name)
            except FileNotFoundError:
                pass


def process_notebook(
    path: Path,
    root: Path,
    coefficients: dict[str, str],
    backup_root: Path,
) -> str:
    if "фитирование" in path.name.casefold():
        print(f"[SKIP fitting] {path}")
        return "skipped_fitting"

    if path.is_symlink():
        print(f"[SKIP symlink] {path}")
        return "skipped_other"

    key = coefficient_key_from_filename(path.name)
    if key is None:
        print(f"[SKIP no-key] {path}: не удалось определить секцию коэффициентов")
        return "skipped_other"

    target = coefficients.get(key)
    if target is None:
        print(f"[SKIP no-coeff] {path}: нет секции {key!r} в coeff_file")
        return "skipped_other"

    try:
        original = path.read_bytes()
    except OSError as exc:
        print(f"[ERROR read] {path}: {exc}")
        return "error"

    try:
        first = find_first_assignment(original)
    except NotebookFormatError as exc:
        print(f"[SKIP first-d] {path}: {exc}")
        return "skipped_first_d"

    if first.variable != "d":
        print(
            f"[SKIP first-d] {path}: d не первая переменная "
            f"(первая переменная: {first.variable})"
        )
        return "skipped_first_d"

    if decimal_equal(first.value_text, target):
        print(f"[OK already] {path}: d = {first.value_text}")
        return "already"

    try:
        patched = patch_d_preserving_size(original, first, target)
    except NotebookFormatError as exc:
        print(f"[ERROR patch] {path}: {exc}")
        return "error"

    try:
        backup_path = backup_file(path, root, backup_root)
    except OSError as exc:
        print(f"[ERROR backup] {path}: {exc}")
        return "error"

    try:
        atomic_write(path, patched)

        # Проверяем фактически записанные байты, а не только буфер.
        written = path.read_bytes()
        if written != patched:
            raise OSError("байты после записи не совпадают с проверенным буфером")

        verified = find_first_assignment(written)
        if verified.variable != "d" or verified.value_text != target:
            raise OSError("d после записи не прошёл контрольную проверку")

    except Exception as exc:
        # Если что-то пошло не так после backup, откатываем исходник.
        try:
            shutil.copy2(backup_path, path)
            restore_note = "исходный файл восстановлен из backup"
        except OSError as restore_exc:
            restore_note = f"ОШИБКА ВОССТАНОВЛЕНИЯ: {restore_exc}"

        print(f"[ERROR write] {path}: {exc}; {restore_note}")
        return "error"

    print(
        f"[CHANGED] {path}: d {first.value_text} -> {target} "
        f"(размер файла сохранён: {len(original)} байт)"
    )
    return "changed"


def main() -> int:
    args = parse_args()

    coeff_file = args.coeff_file.expanduser().resolve()
    root = args.wolfram_notebooks.expanduser().resolve()

    if not coeff_file.is_file():
        print(f"ERROR: coeff_file не является файлом: {coeff_file}", file=sys.stderr)
        return 2

    if not root.is_dir():
        print(
            f"ERROR: wolfram_notebooks не является директорией: {root}",
            file=sys.stderr,
        )
        return 2

    try:
        coefficients = parse_coefficients(coeff_file)
    except RuntimeError as exc:
        print(f"ERROR: {exc}", file=sys.stderr)
        return 2

    # Папка специально НЕ удаляется после завершения.
    backup_root = Path(
        tempfile.mkdtemp(prefix="insert_d_into_nb_backup_", dir="/tmp")
    )

    counters = {
        "changed": 0,
        "already": 0,
        "skipped_fitting": 0,
        "skipped_first_d": 0,
        "skipped_other": 0,
        "error": 0,
    }

    notebooks = sorted(
        (p for p in root.rglob("*.nb") if p.is_file() or p.is_symlink()),
        key=lambda p: str(p),
    )

    print(f"Найдено notebook-файлов: {len(notebooks)}")

    for path in notebooks:
        result = process_notebook(path, root, coefficients, backup_root)
        counters[result] += 1

    print("\nИтог:")
    print(f"  изменено:                         {counters['changed']}")
    print(f"  уже имели правильное d:           {counters['already']}")
    print(f"  пропущено (фитирование):          {counters['skipped_fitting']}")
    print(f"  пропущено (d не первая):          {counters['skipped_first_d']}")
    print(f"  пропущено (прочее):               {counters['skipped_other']}")
    print(f"  ошибок:                           {counters['error']}")
    print(f"Папка backup: {backup_root}")

    return 1 if counters["error"] else 0


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