#!/usr/bin/env python3
# SPDX-FileCopyrightText: 2025 Blender Authors
#
# SPDX-License-Identifier: GPL-2.0-or-later

"""Self-bootstrapping script to run the OpenAPI-to-dataclasses code generator.

Run this via `<your buildtool> generate_datamodels` in your build directory.

This script creates its own virtualenv, installs its dependencies, and then runs
the code generator. It processes OpenAPI spec files in YAML format (see
`YAML_PATHS` below) to generate Python source files. Each `xxx.yaml` file will
produce an `xxx.py` file in the same directory. These Python files also include
the OpenAPI spec, as a Python dict.

The generated Python files are tracked by Git. This generator is NOT part of the
regular Blender build process, and only needs to be run when any of the YAML
files change.
"""

__all__ = (
    "main",
)

import argparse
from pathlib import Path
import sys
import time

# Paths of the OpenAPI YAML files to convert to Python code. These are relative
# to Blender's top level source directory.
#
# The generated Python files will be written to the same path, just with the
# `.py` suffix.
#
# When adding a file here, make sure it is named `..._openapi.yaml`. That way
# the corresponding `.py`` file is automatically marked as 'generated' in
# `.gitattributes`.
YAML_PATHS = [
    "scripts/modules/_bpy_internal/assets/remote_library/blender_asset_library_openapi.yaml",
]

# Packages to install in the virtualenv. These are only necessary to run this
# generator. The generated code does not depend on these.
REQUIREMENTS = [
    "datamodel-code-generator ~= 0.53.0",
    "PyYAML ~= 6.0.2",
    "docformatter ~= 1.7.8",
]

# These arguments are quite likely to be used for all code generated with this
# generator, also later when we use this approach in other areas.
COMMON_ARGS = [
    # Because of the Blender code standard:
    "--use-double-quotes",

    # Make it strict unless there's a good reason not to:
    "--strict-nullable",

    # Ignore unknown fields in the parsed JSON. This way, the code generated now
    # has a chance of being be compatible with future versions of the schema (at
    # least, when that future version just adds new stuff).
    "--allow-extra-fields",

    # Automatically target the currently-running version of Python:
    f"--target-python-version={sys.version_info.major}.{sys.version_info.minor}",

    # Use `list[T]` instead of `typing.List[T]`:
    "--use-standard-collections",

    # Because we use dataclasses:
    "--output-model-type", "dataclasses.dataclass",

    # Work around https://github.com/koxudaxi/datamodel-code-generator/issues/1870#issuecomment-2775689249
    "--use-annotated",

    # Remove the "generated on" timestamp from the output, so that running the
    # generator is idempotent.
    "--disable-timestamp",

    "--use-inline-field-description",
    "--use-schema-description",
]


CUSTOM_FILE_HEADER = """
# SPDX-FileCopyrightText: {year!s} Blender Authors
#
# SPDX-License-Identifier: GPL-2.0-or-later
#
# Generated by datamodel-codegen:
#   source filename:  {source_path.name!s}
"""


def main() -> None:
    """Run the datamodel code generator."""

    argparser = argparse.ArgumentParser(description="Run the datamodel code generator.")
    argparser.add_argument('source_root', type=Path, help="The root of Blender's source directory")
    args = argparser.parse_args(sys.argv[1:])
    root_path: Path = args.source_root.resolve()

    if not root_path.is_dir():
        raise SystemExit("Path {!s} should be a directory".format(root_path))

    print("Generating data model files:")

    py_paths: list[Path] = []
    for yaml_relpath in YAML_PATHS:
        yaml_path = root_path / yaml_relpath
        py_path = yaml_path.with_suffix(".py")
        py_paths.append(py_path)

        print(f"  {yaml_path.relative_to(root_path)} -> {py_path.name}")

        _generate_datamodel(
            in_path=yaml_path,
            in_type="openapi",
            out_path=py_path,
        )

    # Make sure that output from subprocesses is flushed, before outputting more
    # below. This prevents stderr and stdout going out of sync, ensuring things
    # are shown in chronological order (i.e. generating files before
    # reformatting them).
    sys.stderr.flush()
    sys.stdout.flush()

    # Format the generated Python code.
    _docformatter(py_paths)
    _make_format(root_path, py_paths)

    print("Done generating data model files!")


def _generate_datamodel(in_path: Path, in_type: str, out_path: Path) -> None:
    """Run datamodel-codegen."""

    # `type: ignore` to ignore warnings that this module cannot be imported. Python checkers
    # won't understand this is run from a self-managed virtualenv.
    from datamodel_code_generator.__main__ import main as codegen_main  # type: ignore
    from datamodel_code_generator.__main__ import Exit  # type: ignore

    header = CUSTOM_FILE_HEADER.strip().format(
        year=time.localtime().tm_year,
        source_path=in_path,
    )

    args = [
        *COMMON_ARGS,
        "--input", str(in_path),
        "--input-file-type", in_type,
        "--output", str(out_path),
        "--custom-file-header", header,
    ]

    status = codegen_main(args)

    match status:
        case Exit.OK:
            return
        case Exit.ERROR:
            raise SystemExit("code generation failed")
        case Exit.KeyboardInterrupt:
            raise KeyboardInterrupt()
        case _:
            raise SystemExit(f"unknown result from code generation: {status}")


def _docformatter(py_paths: list[Path]) -> None:
    """Run 'docformatter' on generated Python files.

    This is necessary because the generated docstrings are very long, and
    'make format' doesn't automatically re-wrap them.
    """
    from docformatter import format
    from docformatter import configuration

    print("Formatting docstrings")

    argv = ["docformatter", "--in-place", *(str(path) for path in py_paths)]
    cfg = configuration.Configurater(argv)
    cfg.do_parse_arguments()

    formatter = format.Formatter(
        cfg.args,
        stderror=sys.stderr,
        stdin=sys.stdin,
        stdout=sys.stdout,
    )
    result = formatter.do_format_files()
    if result not in {format.FormatResult.ok, format.FormatResult.format_required}:
        raise RuntimeError(f"Error {result} running docformatter")


def _make_format(root_path: Path, py_paths: list[Path]) -> None:
    """Run 'make format' on generated Python files."""

    print("Formatting Python files")

    py_paths_as_str = [str(path) for path in py_paths]
    subprocess.run(
        ["make", "format", "PATHS={}".format(" ".join(py_paths_as_str))],
        cwd=root_path,
        check=True,
    )


# --------- Below this point is the self-bootstrapping logic ---------


import importlib.util
import subprocess
import venv

# Name of a module to import, to test whether dependencies have been installed or not.
TEST_INSTALL_MODULE = "datamodel_code_generator"

# Directory for the virtualenv. This script is expected to run with Blender's
# build directory as its working directory.
VENV_DIR = Path("generate_datamodels_venv").resolve()

# Python executable inside the virtual environment.
VENV_PYTHON = VENV_DIR / "Scripts/python.exe" if sys.platform == "win32" else VENV_DIR / "bin/python"


def _create_virtualenv() -> None:
    """Create the virtual environment if it does not exist."""
    if VENV_DIR.exists():
        return
    print(f"Creating virtual environment at {VENV_DIR}")
    venv.create(VENV_DIR, with_pip=True)


def _install_dependencies() -> None:
    """Install required dependencies into the virtual environment."""
    print("Installing dependencies")
    # Pip doesn't like to be used as Python library, invoking it via the CLI is the best option.
    _run_command(str(VENV_PYTHON), "-m", "pip", "install", "--upgrade", "pip")
    _run_command(str(VENV_PYTHON), "-m", "pip", "install", "--upgrade", *REQUIREMENTS)


def _is_dependency_installed(package: str) -> bool:
    """Try importing a package to check if it is installed."""
    return importlib.util.find_spec(package) is not None


def _is_running_in_virtualenv() -> bool:
    """Check if the script is running inside the virtual environment."""
    return sys.prefix != sys.base_prefix  # Virtualenv modifies `sys.prefix`


def _run_command(*cmd: str) -> None:
    """Run a shell command and handle errors."""
    try:
        subprocess.run(cmd, check=True, text=True)
    except subprocess.CalledProcessError as e:
        print(f"Error running command: {' '.join(cmd)}", file=sys.stderr)
        print(f"Exit code: {e.returncode}", file=sys.stderr)
        sys.exit(e.returncode)


if __name__ == "__main__":
    _create_virtualenv()

    if not _is_running_in_virtualenv():
        print(f"Re-executing inside virtual environment at {VENV_DIR}")
        _run_command(str(VENV_PYTHON), *sys.argv)
        sys.exit()

    if not _is_dependency_installed(TEST_INSTALL_MODULE):
        _install_dependencies()

    # The virtual environment is active, so run the main script logic.
    main()