# Copyright 2026 The HuggingFace Inc. team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Utility that generates `src/transformers/models/__init__.py` and all `src/transformers/models/<model>/__init__.py` files.

Usage (from the root of the repo):

Check that the inits are up to date (used in `make check-repo`):

```bash
python utils/check_inits.py
```

Regenerate them if needed (used in `make fix-repo`):

```bash
python utils/check_inits.py --fix_and_overwrite
```
"""

import argparse
import difflib
import re
from pathlib import Path

from transformers.utils.import_utils import define_import_structure


CHECKER_CONFIG = {
    "name": "inits",
    "label": "Model init files",
    "cache_globs": ["src/transformers/models/**/*.py", "src/transformers/utils/import_utils.py"],
    "check_args": [],
    "fix_args": ["--fix_and_overwrite"],
}

REPO_ROOT = Path(__file__).parent.parent
MODELS_PATH = REPO_ROOT / "src" / "transformers" / "models"
MODELS_INIT_PATH = MODELS_PATH / "__init__.py"

# Add any directories to ignore here.
# Directories starting with `_` are ignored by default.
IGNORED_DIRECTORIES = {"deprecated"}

AUTO_GENERATED_BANNER = """#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
#           This file was automatically generated from the modules in `{module_path}`.
#           Do NOT edit this file manually as any edits will be overwritten by auto-generation of the file.
#                       A module is picked up once it exports a name through `__all__`.
#                   Regenerate the file with: `python utils/check_inits.py --fix_and_overwrite`
#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
"""

INIT_TEMPLATE = """from typing import TYPE_CHECKING

from {dots}utils import _LazyModule
from {dots}utils.import_utils import define_import_structure


if TYPE_CHECKING:
{imports}else:
    import sys

    _file = globals()["__file__"]
    sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
"""


def natural_sort_key(name: str) -> tuple[str | int, ...]:
    """Sort key matching `ruff`'s import ordering, which compares digit runs numerically."""
    return tuple(int(part) if part.isdigit() else part for part in re.split(r"(\d+)", name))


def get_import_structure() -> dict[str, set[str]]:
    """Map each model directory to the modules it exposes, e.g. `{"bert": {"configuration_bert", ...}, ...}`."""
    import_structure = define_import_structure(str(MODELS_INIT_PATH))
    modules_per_model = {}
    for modules in import_structure.values():
        for dotted_module in modules:
            model_name, _, module_name = dotted_module.partition(".")
            if model_name in IGNORED_DIRECTORIES or model_name.startswith("_"):
                continue
            modules_per_model.setdefault(model_name, set()).add(module_name)
    return modules_per_model


def extract_license_header(path: Path) -> str:
    """Return the leading comment block of a file, without the auto-generated banner."""
    if not path.is_file():
        return ""

    header_lines = []
    for line in path.read_text(encoding="utf-8").splitlines():
        if not line.startswith("#"):
            break
        header_lines.append(line)

    banner_end = max((index for index, line in enumerate(header_lines) if "🚨" in line), default=None)
    if banner_end is not None:
        header_lines = header_lines[banner_end + 1 :]

    return "".join(f"{line}\n" for line in header_lines)


def get_license_header(init_path: Path, module_names: list[str]) -> str:
    """Return the init's own license header, falling back to the one of the modules it exposes for a new init."""
    candidates = [init_path] + [init_path.parent / f"{module_name}.py" for module_name in module_names]
    return next((header for header in map(extract_license_header, candidates) if header), "")


def generate_init(init_path: Path, import_lines: list[str], module_names: list[str]) -> str:
    """Render the expected content of an init file."""
    if not import_lines:
        return ""

    module_path = init_path.parent.relative_to(REPO_ROOT).as_posix()
    banner = AUTO_GENERATED_BANNER.format(module_path=module_path)
    dots = "." * (len(init_path.relative_to(MODELS_PATH).parts) + 1)
    imports = "".join(f"    {line}\n" for line in import_lines)
    license_header = get_license_header(init_path=init_path, module_names=module_names)
    return banner + license_header + INIT_TEMPLATE.format(dots=dots, imports=imports)


def generate_all_inits() -> dict[Path, str]:
    """Map every init under `src/transformers/models` to its expected content."""
    modules_per_model = get_import_structure()

    model_names = set(modules_per_model)
    model_names.update(
        path.name for path in MODELS_PATH.iterdir() if path.is_dir() and (path / "__init__.py").is_file()
    )
    model_names -= IGNORED_DIRECTORIES
    model_names = {model_name for model_name in model_names if not model_name.startswith("_")}

    expected_contents = {}
    for model_name in model_names:
        modules = sorted(modules_per_model.get(model_name, set()), key=natural_sort_key)
        import_lines = [f"from .{module_name} import *" for module_name in modules]
        expected_contents[MODELS_PATH / model_name / "__init__.py"] = generate_init(
            init_path=MODELS_PATH / model_name / "__init__.py", import_lines=import_lines, module_names=modules
        )

    expected_contents[MODELS_INIT_PATH] = generate_init(
        init_path=MODELS_INIT_PATH,
        module_names=[],
        import_lines=[
            f"from .{model_name} import *" for model_name in sorted(modules_per_model, key=natural_sort_key)
        ],
    )
    return expected_contents


def main(overwrite: bool):
    diffs = []
    for init_path, new_content in sorted(generate_all_inits().items()):
        old_content = init_path.read_text(encoding="utf-8") if init_path.is_file() else ""
        if old_content == new_content:
            continue

        if overwrite:
            init_path.parent.mkdir(parents=True, exist_ok=True)
            init_path.write_text(new_content, encoding="utf-8")
            continue

        relative_path = init_path.relative_to(REPO_ROOT)
        diffs.append(
            "".join(
                difflib.unified_diff(
                    old_content.splitlines(keepends=True),
                    new_content.splitlines(keepends=True),
                    fromfile=f"{relative_path} (on disk)",
                    tofile=f"{relative_path} (regenerated)",
                )
            )
        )

    if diffs:
        raise Exception(
            f"{len(diffs)} init file(s) are not consistent with the import structure on disk.\n"
            "Run `make fix-repo` or `python utils/check_inits.py --fix_and_overwrite` to fix them.\n\n"
            "Diff (on disk → regenerated):\n" + "\n".join(diffs)
        )


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--fix_and_overwrite", action="store_true", help="Whether to fix inconsistencies.")
    args = parser.parse_args()
    main(overwrite=args.fix_and_overwrite)