#!/usr/bin/env python3
"""Prepare a temporary source tree for docs generation.

The live source keeps readable display-math blocks written as standalone
`$$ ... $$` docstring sections. MkDocs + mkdocstrings does not consistently
preserve those blocks as MathJax input when they appear inside API docstrings,
so the docs build rewrites them in a temporary copy into explicit
`<div class="arithmatex">...</div>` containers.
"""

from __future__ import annotations

from pathlib import Path
import shutil
import sys


def _rewrite_display_math(text: str) -> str:
    lines = text.splitlines(keepends=True)
    output: list[str] = []
    in_block = False
    block_indent = ""

    for line in lines:
        stripped = line.strip()
        indent = line[: len(line) - len(line.lstrip())]

        if not in_block:
            if stripped == "$$":
                block_indent = indent
                output.append(f'{block_indent}<div class="arithmatex">\\[\n')
                in_block = True
                continue

            if stripped.startswith("$$") and stripped.endswith("$$") and stripped != "$$":
                body = stripped[2:-2].strip()
                output.append(f'{indent}<div class="arithmatex">\\[ {body} \\]</div>\n')
                continue

            if stripped.startswith("$$"):
                block_indent = indent
                body = stripped[2:].strip()
                output.append(f'{block_indent}<div class="arithmatex">\\[\n')
                if body:
                    output.append(f"{block_indent}{body}\n")
                in_block = True
                continue

            output.append(line)
            continue

        if stripped == "$$":
            output.append(f"{block_indent}\\]</div>\n")
            in_block = False
            block_indent = ""
            continue

        if stripped.endswith("$$"):
            body = stripped[:-2].rstrip()
            if body:
                output.append(f"{block_indent}{body}\n")
            output.append(f"{block_indent}\\]</div>\n")
            in_block = False
            block_indent = ""
            continue

        output.append(line)

    if in_block:
        raise ValueError("unterminated display-math block")

    return "".join(output)


def main() -> int:
    if len(sys.argv) != 3:
        print("usage: prepare_docs_sources.py <src_package_dir> <dst_root>", file=sys.stderr)
        return 2

    src_dir = Path(sys.argv[1]).resolve()
    dst_root = Path(sys.argv[2]).resolve()
    dst_pkg = dst_root / src_dir.name

    shutil.copytree(src_dir, dst_pkg)

    for path in dst_pkg.rglob("*.py"):
        path.write_text(_rewrite_display_math(path.read_text()))

    return 0


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