"""Point deepinfer's `tile_kernels` imports at the compat layer.

Two rewrites, and only on files the model forward imports:

    import tile_kernels                     -> from deepinfer import tk_compat as tile_kernels
    from tile_kernels.<mod> import <names>  -> from deepinfer.tk_compat.<mod> import <names>

`from tile_kernels.<mod>.<submod> import x` (e.g. `tile_kernels.mhc.post_kernel`) is rewritten to
`from deepinfer.tk_compat.<mod> import x`, because the compat layer flattens the per-implementation
submodules into one module per group.

Files under `tilelang_ops/asc/` and `tilelang_ops/cuda/` are left alone: those are backend
implementations that only ever load when TileLang is present, and routing them through the
reference path would be wrong rather than merely unnecessary.

    python replace_imports.py --check    # report what would change
    python replace_imports.py           # rewrite
    python replace_imports.py --revert   # restore from .import-orig backups
"""

import argparse
import re
import shutil
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[2] / "deepinfer"
# Backend implementations only load when TileLang is present, so routing them through the
# reference path would be wrong rather than merely unnecessary. Matched as path segments because
# the files sit several directories deep.
SKIP_PARTS = {("tilelang_ops", "asc"), ("tilelang_ops", "cuda"), ("tk_compat",)}

BARE = "import tile_kernels\n"
BARE_NEW = "from deepinfer import tk_compat as tile_kernels\n"

# Anchored to the start of the line, i.e. module-level only. A function-local
# `import tile_kernels` is how the tilelang_ops reference implementations load the *compiled*
# kernels to compare their Triton ports against; routing those through the reference path would
# compare a thing to itself.
BARE_RE = re.compile(r"^import tile_kernels$", re.MULTILINE)

# from tile_kernels.a.b import x,y  /  from tile_kernels.a import x
FROM_RE = re.compile(r"^from (tile_kernels(?:\.\w+)+) import (.+)$", re.MULTILINE)


def targeted_files() -> list[Path]:
    out = []
    for path in sorted(ROOT.rglob("*.py")):
        parts = path.relative_to(ROOT).parts
        if any(all(seg in parts for seg in group) for group in SKIP_PARTS):
            continue
        text = path.read_text()
        if BARE_RE.search(text) or re.search(r"^from tile_kernels", text, re.MULTILINE):
            out.append(path)
    return out


def rewrite(text: str) -> str:
    text = BARE_RE.sub(BARE_NEW.rstrip("\n"), text)

    def sub(match: re.Match[str]) -> str:
        module, names = match.group(1), match.group(2)
        parts = module.split(".")[1:]  # drop the leading `tile_kernels`
        # `tile_kernels.quant` -> quant; `tile_kernels.mhc.post_kernel` -> mhc
        group = parts[0]
        return f"from deepinfer.tk_compat.{group} import {names}"

    return FROM_RE.sub(sub, text)


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

    files = targeted_files()
    if args.revert:
        restored = 0
        for path in files:
            backup = path.with_suffix(".py.import-orig")
            if backup.is_file():
                shutil.move(str(backup), str(path))
                restored += 1
        print(f"restored {restored} file(s)")
        return 0

    changed = 0
    for path in files:
        original = path.read_text()
        updated = rewrite(original)
        if updated == original:
            continue
        rel = path.relative_to(ROOT.parent).as_posix()
        if args.check:
            print(f"would rewrite {rel}")
        else:
            backup = path.with_suffix(".py.import-orig")
            if not backup.is_file():
                shutil.copy2(path, backup)
            path.write_text(updated)
            print(f"rewrote {rel}")
        changed += 1

    print(f"\n{changed} file(s) {'would change' if args.check else 'changed'}")
    return 0


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