"""Rewrite an eagerly-importing package `__init__.py` into a lazy one (PEP 562).

Some TileLang packages re-export every compiled kernel from their `__init__`. That makes the
package unimportable without TileLang even when the caller only wants a pure-Python submodule
-- importing `pkg.sub` runs `pkg/__init__.py` first. Resolving the re-exports on first access
fixes it without changing what callers see.

The export map is parsed from the original file rather than hand-written, so a name this script
does not know about cannot silently disappear.

`__all__` is emitted so `from pkg import *` keeps working: wildcard imports call `getattr` for
every name listed there, which is exactly what triggers the lazy resolution.

    python lazify_package_init.py <path/to/__init__.py> [--revert]
"""

import argparse
import ast
import shutil
import sys
from pathlib import Path

TEMPLATE = '''"""Lazy re-exports, so importing a submodule does not import every compiled kernel.

Generated from the original eager `__init__.py`: each name below is resolved from its module on
first access. The reason is that `import {pkg}.sub` executes this file first, so an eager
re-export list made even the pure-Python submodules unavailable on a host without TileLang.
"""

import importlib

_EXPORTS = {{
{exports}
}}

__all__ = sorted(_EXPORTS)


def __getattr__(name: str):
    module = _EXPORTS.get(name)
    if module is None:
        raise AttributeError(f"module {{__name__!r}} has no attribute {{name!r}}")
    value = getattr(importlib.import_module(module, __name__), name)
    globals()[name] = value
    return value


def __dir__() -> list[str]:
    return sorted(set(globals()) | set(_EXPORTS))
'''


def collect(tree: ast.Module) -> dict[str, str]:
    """Map each re-exported name to the module it comes from, as an absolute import string."""
    exports: dict[str, str] = {}
    for node in tree.body:
        if isinstance(node, ast.ImportFrom) and node.level == 1 and node.module:
            # `from .mod import a, b`
            for alias in node.names:
                if alias.name == "*":
                    raise SystemExit("wildcard import in __init__; cannot infer the export map")
                exports[alias.asname or alias.name] = f".{node.module}"
        elif isinstance(node, ast.ImportFrom) and node.level == 1 and node.module is None:
            # `from . import mod` -- the export *is* the submodule.
            for alias in node.names:
                exports[alias.asname or alias.name] = f".{alias.name}"
        elif isinstance(node, (ast.Import, ast.FunctionDef, ast.ClassDef, ast.Assign, ast.AnnAssign)):
            raise SystemExit(f"unexpected top-level statement in __init__: {type(node).__name__}")
    return exports


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("init_py", type=Path)
    parser.add_argument("--revert", action="store_true")
    args = parser.parse_args()

    path = args.init_py
    backup = path.with_suffix(".py.orig")
    if not path.is_file():
        print(f"missing: {path}", file=sys.stderr)
        return 2

    if args.revert:
        if not backup.is_file():
            print(f"no backup: {backup}", file=sys.stderr)
            return 1
        shutil.move(str(backup), str(path))
        print(f"reverted {path}")
        return 0

    source = path.read_text()
    if "_EXPORTS" in source:
        print("already lazy; nothing to do")
        return 0
    exports = collect(ast.parse(source))
    if not exports:
        print("no re-exports found; nothing to do")
        return 0

    if not backup.is_file():
        shutil.copy2(path, backup)

    width = max(len(name) for name in exports)
    body = "\n".join(f'    "{name}":'.ljust(width + 6) + f' "{module}",' for name, module in sorted(exports.items()))
    path.write_text(TEMPLATE.format(pkg=path.parent.name, exports=body))
    compile(path.read_text(), str(path), "exec")
    print(f"lazified {path}: {len(exports)} names")
    for name, module in sorted(exports.items()):
        print(f"  {name:<45} {module}")
    return 0


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