#!/usr/bin/env python3
"""Удаляет из текстовых Wolfram Notebook (.nb) ячейки с Manipulate/CellPrint.

Удаляется вся вычислительная группа CellGroupData, если одна из её Input-ячеек
содержит Manipulate[...] или CellPrint[...]. Поэтому вместе с исходной ячейкой
исчезают связанные Output/Message/Print-ячейки.

Служебные блоки Notebook (CacheID, NotebookFileOutline, CellTagsIndex и весь
Internal cache information) не анализируются и не изменяются. Все нетронутые
части списка ячеек также сохраняются посимвольно.

Примеры:
    python clean_wolfram_notebook.py notebook.nb
    python clean_wolfram_notebook.py a.nb b.nb --output-dir cleaned
    python clean_wolfram_notebook.py notebook.nb --in-place
    python clean_wolfram_notebook.py notebook.nb --dry-run
"""

from __future__ import annotations

import argparse
import os
import re
import shutil
import sys
import tempfile
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable, Optional


RESULT_STYLES = {
    "Output",
    "Print",
    "Message",
    "Echo",
    "GeneratedCell",
}


class NotebookFormatError(RuntimeError):
    """Файл не похож на корректный текстовый Wolfram Notebook."""


@dataclass
class Stats:
    manipulate_groups: int = 0
    cellprint_groups: int = 0
    manipulate_cells: int = 0
    cellprint_cells: int = 0
    result_cells: int = 0

    @property
    def removed_groups(self) -> int:
        return self.manipulate_groups + self.cellprint_groups

    @property
    def removed_input_cells(self) -> int:
        return self.manipulate_cells + self.cellprint_cells

    def add(self, other: "Stats") -> None:
        self.manipulate_groups += other.manipulate_groups
        self.cellprint_groups += other.cellprint_groups
        self.manipulate_cells += other.manipulate_cells
        self.cellprint_cells += other.cellprint_cells
        self.result_cells += other.result_cells


@dataclass(frozen=True)
class ParsedCall:
    open_index: int
    close_index: int
    args: list[str]
    arg_ranges: list[tuple[int, int]]


@dataclass(frozen=True)
class ParsedList:
    open_index: int
    close_index: int
    items: list[str]
    item_ranges: list[tuple[int, int]]


@dataclass(frozen=True)
class InputTarget:
    manipulate: bool = False
    cellprint: bool = False

    @property
    def any(self) -> bool:
        return self.manipulate or self.cellprint


def _scan_top_level_ranges(text: str, start: int, end: int) -> list[tuple[int, int]]:
    """Разбивает диапазон по запятым верхнего уровня.

    Учитываются строки, экранирование, вложенные (), [] и {}, а также
    вложенные комментарии Wolfram вида (* ... *).
    """
    ranges: list[tuple[int, int]] = []
    item_start = start
    stack: list[str] = []
    pairs = {")": "(", "]": "[", "}": "{"}
    in_string = False
    escaped = False
    comment_depth = 0
    i = start

    while i < end:
        ch = text[i]
        nxt = text[i + 1] if i + 1 < end else ""

        if comment_depth:
            if ch == "(" and nxt == "*":
                comment_depth += 1
                i += 2
                continue
            if ch == "*" and nxt == ")":
                comment_depth -= 1
                i += 2
                continue
            i += 1
            continue

        if in_string:
            if escaped:
                escaped = False
            elif ch == "\\":
                escaped = True
            elif ch == '"':
                in_string = False
            i += 1
            continue

        if ch == "(" and nxt == "*":
            comment_depth = 1
            i += 2
            continue
        if ch == '"':
            in_string = True
            i += 1
            continue
        if ch in "([{":
            stack.append(ch)
            i += 1
            continue
        if ch in ")]}" :
            if not stack or stack[-1] != pairs[ch]:
                raise NotebookFormatError(
                    f"Несогласованная скобка {ch!r} в позиции {i}"
                )
            stack.pop()
            i += 1
            continue
        if ch == "," and not stack:
            ranges.append((item_start, i))
            item_start = i + 1
        i += 1

    if in_string:
        raise NotebookFormatError("Незакрытая строка в Notebook")
    if comment_depth:
        raise NotebookFormatError("Незакрытый комментарий в Notebook")
    if stack:
        raise NotebookFormatError("Незакрытая скобка в Notebook")

    ranges.append((item_start, end))
    return ranges


def _find_matching(text: str, open_index: int) -> int:
    opener = text[open_index]
    closer = {"[": "]", "{": "}", "(": ")"}.get(opener)
    if closer is None:
        raise ValueError(f"Ожидалась открывающая скобка, получено {opener!r}")

    stack: list[str] = [opener]
    pairs = {")": "(", "]": "[", "}": "{"}
    in_string = False
    escaped = False
    comment_depth = 0
    i = open_index + 1

    while i < len(text):
        ch = text[i]
        nxt = text[i + 1] if i + 1 < len(text) else ""

        if comment_depth:
            if ch == "(" and nxt == "*":
                comment_depth += 1
                i += 2
                continue
            if ch == "*" and nxt == ")":
                comment_depth -= 1
                i += 2
                continue
            i += 1
            continue

        if in_string:
            if escaped:
                escaped = False
            elif ch == "\\":
                escaped = True
            elif ch == '"':
                in_string = False
            i += 1
            continue

        if ch == "(" and nxt == "*":
            comment_depth = 1
            i += 2
            continue
        if ch == '"':
            in_string = True
            i += 1
            continue
        if ch in "([{":
            stack.append(ch)
            i += 1
            continue
        if ch in ")]}" :
            if not stack or stack[-1] != pairs[ch]:
                raise NotebookFormatError(
                    f"Несогласованная скобка {ch!r} в позиции {i}"
                )
            stack.pop()
            if not stack:
                return i
            i += 1
            continue
        i += 1

    raise NotebookFormatError(
        f"Не найдена закрывающая скобка {closer!r} для позиции {open_index}"
    )


def _skip_space_and_comments(text: str, pos: int, end: Optional[int] = None) -> int:
    limit = len(text) if end is None else end
    while pos < limit:
        if text[pos].isspace():
            pos += 1
            continue
        if pos + 1 < limit and text[pos : pos + 2] == "(*":
            depth = 1
            pos += 2
            while pos < limit and depth:
                if pos + 1 < limit and text[pos : pos + 2] == "(*":
                    depth += 1
                    pos += 2
                elif pos + 1 < limit and text[pos : pos + 2] == "*)":
                    depth -= 1
                    pos += 2
                else:
                    pos += 1
            if depth:
                raise NotebookFormatError("Незакрытый комментарий")
            continue
        break
    return pos


def _parse_call(expr: str, head: str) -> Optional[ParsedCall]:
    pos = _skip_space_and_comments(expr, 0)
    if not expr.startswith(head, pos):
        return None
    pos += len(head)
    pos = _skip_space_and_comments(expr, pos)
    if pos >= len(expr) or expr[pos] != "[":
        return None
    close = _find_matching(expr, pos)
    tail = _skip_space_and_comments(expr, close + 1)
    if tail != len(expr):
        return None
    ranges = _scan_top_level_ranges(expr, pos + 1, close)
    return ParsedCall(
        open_index=pos,
        close_index=close,
        args=[expr[a:b] for a, b in ranges],
        arg_ranges=ranges,
    )


def _parse_list(expr: str) -> Optional[ParsedList]:
    pos = _skip_space_and_comments(expr, 0)
    if pos >= len(expr) or expr[pos] != "{":
        return None
    close = _find_matching(expr, pos)
    tail = _skip_space_and_comments(expr, close + 1)
    if tail != len(expr):
        return None
    ranges = _scan_top_level_ranges(expr, pos + 1, close)
    # Пустой список даёт один пустой диапазон — нормализуем его.
    if len(ranges) == 1 and not expr[ranges[0][0] : ranges[0][1]].strip():
        ranges = []
    return ParsedList(
        open_index=pos,
        close_index=close,
        items=[expr[a:b] for a, b in ranges],
        item_ranges=ranges,
    )


def _quoted_string(expr: str) -> Optional[str]:
    s = expr.strip()
    if len(s) < 2 or s[0] != '"' or s[-1] != '"':
        return None
    # Для стилей Notebook обычно нет escape-последовательностей.
    return s[1:-1]


def _cell_style(cell_expr: str) -> Optional[str]:
    cell = _parse_call(cell_expr.strip(), "Cell")
    if cell is None or len(cell.args) < 2:
        return None
    return _quoted_string(cell.args[1])


def _input_target(cell_expr: str) -> InputTarget:
    if _cell_style(cell_expr) != "Input":
        return InputTarget()

    # Стандартное коробочное представление: RowBox[{"Manipulate", "[", ...}]
    manipulate = bool(
        re.search(r'"Manipulate"\s*,\s*"\\?\["', cell_expr)
        or re.search(r'(?<![A-Za-z0-9_$`])Manipulate\s*\[', cell_expr)
    )
    cellprint = bool(
        re.search(r'"CellPrint"\s*,\s*"\\?\["', cell_expr)
        or re.search(r'(?<![A-Za-z0-9_$`])CellPrint\s*\[', cell_expr)
    )
    return InputTarget(manipulate=manipulate, cellprint=cellprint)


def _join_modified_items(original_content: str, items: Iterable[str]) -> str:
    """Соединяет оставшиеся элементы, не форматируя их заново.

    Каждый элемент передаётся как точный срез исходного файла вместе с его
    пробелами и переводами строк. Меняются только удалённые элементы и
    разделяющие их запятые.
    """
    kept = [item for item in items if item.strip()]
    if not kept:
        # Пустое содержимое списка допустимо. Намеренно не переносим сюда
        # пробелы удалённых ячеек, чтобы не оставлять огромные пустые области.
        return ""
    return ",".join(kept)


def _clean_cell_group(cell_expr: str) -> tuple[Optional[str], Stats]:
    stats = Stats()
    cell = _parse_call(cell_expr.strip(), "Cell")
    if cell is None or not cell.args:
        return cell_expr, stats

    group_arg = cell.args[0]
    group = _parse_call(group_arg.strip(), "CellGroupData")
    if group is None or not group.args:
        return cell_expr, stats

    children_expr = group.args[0]
    children_stripped = children_expr.strip()
    children = _parse_list(children_stripped)
    if children is None:
        return cell_expr, stats

    # Если непосредственно в этой вычислительной группе есть целевой Input,
    # удаляем всю группу: так гарантированно уходят все её Output/Message/Print.
    direct_target = InputTarget()
    for child in children.items:
        target = _input_target(child.strip())
        direct_target = InputTarget(
            manipulate=direct_target.manipulate or target.manipulate,
            cellprint=direct_target.cellprint or target.cellprint,
        )

    if direct_target.any:
        if direct_target.manipulate:
            stats.manipulate_groups += 1
        elif direct_target.cellprint:
            stats.cellprint_groups += 1
        # В редком случае обе конструкции находятся в одной группе, считаем
        # группу один раз, с приоритетом Manipulate, но ячейки ниже не нужны.
        return None, stats

    cleaned_items, child_stats, changed = _clean_cell_list(
        children_stripped[children.open_index + 1 : children.close_index]
    )
    stats.add(child_stats)
    if not changed:
        return cell_expr, stats

    new_children_stripped = (
        children_stripped[: children.open_index + 1]
        + cleaned_items
        + children_stripped[children.close_index :]
    )
    children_leading = children_expr[: len(children_expr) - len(children_expr.lstrip())]
    children_trailing = children_expr[len(children_expr.rstrip()) :]
    new_children_expr = children_leading + new_children_stripped + children_trailing

    # Заменяем первый аргумент CellGroupData внутри исходного выражения.
    group_stripped = group_arg.strip()
    group_parsed = _parse_call(group_stripped, "CellGroupData")
    assert group_parsed is not None
    arg0_start, arg0_end = group_parsed.arg_ranges[0]
    new_group_stripped = (
        group_stripped[:arg0_start]
        + new_children_expr
        + group_stripped[arg0_end:]
    )
    group_leading = group_arg[: len(group_arg) - len(group_arg.lstrip())]
    group_trailing = group_arg[len(group_arg.rstrip()) :]
    new_group_arg = group_leading + new_group_stripped + group_trailing

    cell_stripped = cell_expr.strip()
    cell_parsed = _parse_call(cell_stripped, "Cell")
    assert cell_parsed is not None
    cell_arg0_start, cell_arg0_end = cell_parsed.arg_ranges[0]
    new_cell_stripped = (
        cell_stripped[:cell_arg0_start]
        + new_group_arg
        + cell_stripped[cell_arg0_end:]
    )
    leading = cell_expr[: len(cell_expr) - len(cell_expr.lstrip())]
    trailing = cell_expr[len(cell_expr.rstrip()) :]
    return leading + new_cell_stripped + trailing, stats


def _clean_cell_list(content: str) -> tuple[str, Stats, bool]:
    ranges = _scan_top_level_ranges(content, 0, len(content))
    if len(ranges) == 1 and not content[ranges[0][0] : ranges[0][1]].strip():
        return content, Stats(), False

    stats = Stats()
    kept: list[str] = []
    changed = False
    remove_following_results = False

    for start, end in ranges:
        raw_item = content[start:end]
        item = raw_item.strip()
        if not item:
            continue

        style = _cell_style(item)

        if remove_following_results:
            if style in RESULT_STYLES:
                stats.result_cells += 1
                changed = True
                continue
            remove_following_results = False

        target = _input_target(item)
        if target.any:
            if target.manipulate:
                stats.manipulate_cells += 1
            elif target.cellprint:
                stats.cellprint_cells += 1
            changed = True
            remove_following_results = True
            continue

        cleaned_group, group_stats = _clean_cell_group(raw_item)
        stats.add(group_stats)
        if cleaned_group is None:
            changed = True
            continue
        if cleaned_group != raw_item:
            changed = True
        kept.append(cleaned_group)

    if not changed:
        return content, stats, False
    return _join_modified_items(content, kept), stats, True


def _find_notebook_cell_list(text: str) -> tuple[int, int]:
    marker = "(* Beginning of Notebook Content *)"
    search_from = text.find(marker)
    if search_from < 0:
        search_from = 0

    match = re.search(r"\bNotebook\s*\[", text[search_from:])
    if not match:
        raise NotebookFormatError("Не найдено выражение Notebook[...]")
    notebook_start = search_from + match.start()
    bracket = text.find("[", notebook_start)
    pos = _skip_space_and_comments(text, bracket + 1)
    if pos >= len(text) or text[pos] != "{":
        raise NotebookFormatError("Первый аргумент Notebook не является списком ячеек")
    close = _find_matching(text, pos)
    return pos, close


def clean_notebook_text(text: str) -> tuple[str, Stats, bool]:
    """Очищает только первый аргумент Notebook[{...}], не трогая кэш.

    Префикс до списка ячеек и суффикс после него вставляются в результат без
    каких-либо преобразований. Это гарантирует сохранность CacheID и всего
    блока Internal cache information.
    """
    list_open, list_close = _find_notebook_cell_list(text)
    prefix = text[: list_open + 1]
    original_content = text[list_open + 1 : list_close]
    suffix = text[list_close:]

    cleaned_content, stats, changed = _clean_cell_list(original_content)
    cleaned = prefix + cleaned_content + suffix if changed else text

    # Проверяем лишь синтаксическую доступность Notebook и списка ячеек.
    # Служебные блоки после Notebook[...] намеренно не читаются и не меняются.
    _find_notebook_cell_list(cleaned)
    return cleaned, stats, changed


def _read_text(path: Path) -> tuple[str, str, bool]:
    data = path.read_bytes()
    has_bom = data.startswith(b"\xef\xbb\xbf")
    if has_bom:
        data = data[3:]
    try:
        text = data.decode("utf-8")
    except UnicodeDecodeError as exc:
        raise NotebookFormatError(
            f"{path}: ожидается текстовый .nb в UTF-8/ASCII: {exc}"
        ) from exc

    newline = "\r\n" if b"\r\n" in data else "\n"
    text = text.replace("\r\n", "\n").replace("\r", "\n")
    return text, newline, has_bom


def _encode_text(text: str, newline: str, has_bom: bool) -> bytes:
    if newline != "\n":
        text = text.replace("\n", newline)
    data = text.encode("utf-8")
    return (b"\xef\xbb\xbf" + data) if has_bom else data


def _atomic_write(path: Path, data: bytes, mode: int = 0o644) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    fd, temp_name = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent)
    temp_path = Path(temp_name)
    try:
        with os.fdopen(fd, "wb") as fh:
            fh.write(data)
            fh.flush()
            os.fsync(fh.fileno())
        os.chmod(temp_path, mode)
        os.replace(temp_path, path)
    except Exception:
        temp_path.unlink(missing_ok=True)
        raise


def _default_output_path(input_path: Path) -> Path:
    return input_path.with_name(f"{input_path.stem}_cleaned{input_path.suffix}")


def _resolve_outputs(args: argparse.Namespace, inputs: list[Path]) -> list[Path]:
    if args.output:
        if len(inputs) != 1:
            raise SystemExit("Параметр --output можно использовать только с одним входным файлом")
        return [Path(args.output)]
    if args.output_dir:
        out_dir = Path(args.output_dir)
        return [out_dir / f"{p.stem}_cleaned{p.suffix}" for p in inputs]
    if args.in_place:
        return inputs
    return [_default_output_path(p) for p in inputs]


def _print_stats(path: Path, stats: Stats, changed: bool, dry_run: bool) -> None:
    action = "Проверен" if dry_run else "Готов"
    status = "изменения найдены" if changed else "целевых ячеек нет"
    print(
        f"{action}: {path} — {status}; "
        f"групп Manipulate: {stats.manipulate_groups}, "
        f"групп CellPrint: {stats.cellprint_groups}, "
        f"отдельных Input: {stats.removed_input_cells}, "
        f"отдельных результатов: {stats.result_cells}"
    )


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description=(
            "Удаляет из Wolfram Notebook группы с Manipulate/CellPrint "
            "вместе с их результатами."
        )
    )
    parser.add_argument("inputs", nargs="+", help="входные .nb-файлы")
    out = parser.add_mutually_exclusive_group()
    out.add_argument("-o", "--output", help="выходной файл (только для одного входа)")
    out.add_argument("--output-dir", help="каталог для очищенных файлов")
    out.add_argument(
        "--in-place",
        action="store_true",
        help="перезаписать входные файлы; рядом создаётся резервная копия .bak",
    )
    parser.add_argument(
        "--no-backup",
        action="store_true",
        help="не создавать .bak при --in-place",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="только показать, что будет удалено, не записывая файлы",
    )
    return parser


def main(argv: Optional[list[str]] = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)
    inputs = [Path(p) for p in args.inputs]

    for path in inputs:
        if not path.is_file():
            parser.error(f"файл не найден: {path}")
        if path.suffix.lower() != ".nb":
            parser.error(f"ожидался файл .nb: {path}")

    outputs = _resolve_outputs(args, inputs)
    exit_code = 0

    for input_path, output_path in zip(inputs, outputs):
        try:
            text, newline, has_bom = _read_text(input_path)
            cleaned, stats, changed = clean_notebook_text(text)

            if args.dry_run:
                _print_stats(input_path, stats, changed, dry_run=True)
                continue

            if args.in_place and changed and not args.no_backup:
                backup = input_path.with_suffix(input_path.suffix + ".bak")
                shutil.copy2(input_path, backup)

            input_mode = input_path.stat().st_mode & 0o777
            _atomic_write(
                output_path,
                _encode_text(cleaned, newline, has_bom),
                mode=input_mode or 0o644,
            )
            _print_stats(output_path, stats, changed, dry_run=False)
        except (OSError, NotebookFormatError) as exc:
            print(f"Ошибка: {input_path}: {exc}", file=sys.stderr)
            exit_code = 1

    return exit_code


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