# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""Run stubgen-pyx for cuda_core and normalize the generated stub headers.

stubgen-pyx emits a path using the OS path separator in the first-line comment
(e.g. "# This file was generated by stubgen-pyx from cuda_core\\cuda\\...").
This wrapper rewrites that separator to "/" so committed stubs are identical
across platforms. Line-ending normalization is handled by .gitattributes.

This also forces stubgen-pyx to write files with UTF-8 encoding, which is not
the default on Windows.

This wrapper can be removed once these stubgen-pyx issues are resolved:
    https://github.com/jon-edward/stubgen-pyx/issues/41
    https://github.com/jon-edward/stubgen-pyx/issues/42
"""

from __future__ import annotations

import os
import pathlib
import subprocess
import sys

_HEADER_PREFIX = b"# This file was generated by stubgen-pyx"


def _normalize_stub_headers(root: pathlib.Path) -> None:
    for stub in root.rglob("*.pyi"):
        data = stub.read_bytes()
        newline = data.find(b"\n")
        first_line = data[:newline] if newline != -1 else data
        if not first_line.startswith(_HEADER_PREFIX) or b"\\" not in first_line:
            continue
        stub.write_bytes(first_line.replace(b"\\", b"/") + data[newline:])


def main() -> int:
    env = os.environ.copy()
    env.setdefault("PYTHONUTF8", "1")
    env.setdefault("PYTHONIOENCODING", "utf-8")
    result = subprocess.run(
        ["stubgen-pyx", "cuda_core/cuda", "--continue-on-error", "--include-private"],  # noqa: S607
        env=env,
    )
    if result.returncode != 0:
        return result.returncode
    _normalize_stub_headers(pathlib.Path("cuda_core/cuda"))
    return 0


if __name__ == "__main__":
    sys.exit(main())
