"""Audit the public source and distributions for the inference-only boundary."""

from __future__ import annotations

import argparse
import ast
import tarfile
import zipfile
from pathlib import Path

RUNTIME_MODULES = {
    "__init__",
    "__main__",
    "benchmark",
    "cli",
    "compare",
    "config",
    "dit",
    "encoder",
    "hub",
    "kernels",
    "layout",
    "memory",
    "ops",
    "packed",
    "reference",
    "runtime",
    "tuning",
    "vae",
    "vision",
    "weights",
}


def check_python(name, source):
    tree = ast.parse(source, filename=name)
    for node in ast.walk(tree):
        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
            if node.name in {"pack", "from_linear", "convert", "export_model"}:
                raise ValueError(f"Conversion entry point in {name}: {node.name}")
        if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute):
            if node.func.attr in {"quantize", "save_model", "export_model"}:
                raise ValueError(f"Conversion operation in {name}: {node.func.attr}")
        if isinstance(node, (ast.Import, ast.ImportFrom)):
            names = [a.name for a in node.names]
            if isinstance(node, ast.ImportFrom) and node.module:
                names.append(node.module)
            if any(
                n.split(".")[0] in {"experiments", "scripts"} or n.endswith(".convert")
                for n in names
            ):
                raise ValueError(f"Private tooling import in {name}")


def check_entry(name, data):
    path = Path(name)
    if path.suffix in {".safetensors", ".gguf", ".npz", ".npy"}:
        raise ValueError(f"Model or tensor artifact in public distribution: {name}")
    if path.suffix == ".py" and "h3_mlx" in path.parts:
        if path.stem not in RUNTIME_MODULES:
            raise ValueError(f"Unreviewed runtime module: {name}")
        check_python(name, data)


def check_source(root):
    for path in (root / "h3_mlx").glob("*.py"):
        check_entry(str(path), path.read_bytes())
    for name in ("experiments", "optimization", ".release"):
        if (root / name).exists():
            raise ValueError(
                f"Private development directory in public checkout: {name}"
            )


def check_artifact(path):
    if path.suffix == ".whl":
        with zipfile.ZipFile(path) as archive:
            for name in archive.namelist():
                check_entry(name, archive.read(name))
    elif path.name.endswith(".tar.gz"):
        with tarfile.open(path) as archive:
            for member in archive.getmembers():
                if member.isfile():
                    check_entry(member.name, archive.extractfile(member).read())
    else:
        raise ValueError(f"Unsupported distribution: {path}")


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("artifacts", nargs="*", type=Path)
    args = parser.parse_args()
    check_source(Path(__file__).resolve().parents[1])
    for path in args.artifacts:
        check_artifact(path)
    print(
        f"Inference-only audit passed: source and {len(args.artifacts)} distributions"
    )


if __name__ == "__main__":
    main()