"""Copy the game files SUPERBEAM needs into envs/superbeam_env/vendor/ so the env dir
(and its Docker image / HF Space) is self-contained.  Re-run after changing the game.

  python scripts/prepare_superbeam_env.py

The imports of beamboot/superbeam_gym.py are followed automatically (module-level
imports of beamboot.* only). The in-weights computer (chip8_lm, chip8_fast), the scripts
that run and inspect it, and its tests are copied too, so the Hub repo holds the full
source; they need torch + transformers (Beam), which the env server image does not
install (superbeam_gym imports chip8_fast lazily, only for the beam-weights backend).
Run its tests from vendor/:  cd vendor && pytest tests/
"""
import ast
import os
import shutil

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
DEST = os.path.join(ROOT, "envs", "superbeam_env", "vendor")
ROMS = ("superbeam.asm", "superbeam.ch8", "superbeam_inputs.json")
IN_WEIGHTS = ("beamboot/chip8_lm.py", "beamboot/chip8_fast.py",
              "scripts/run_chip8_fast.py", "scripts/attention_probe.py", "scripts/deploy_inweights_modal.py",
              "tests/test_chip8_lm.py", "tests/test_chip8_fast.py")


def module_level_deps(relpath, seen):
    if relpath in seen:
        return
    seen.add(relpath)
    tree = ast.parse(open(os.path.join(ROOT, relpath)).read())
    for node in tree.body:                                  # module level only: lazy imports are skipped
        names = []
        if isinstance(node, ast.ImportFrom) and node.module:
            names = [node.module] + [node.module + "." + a.name for a in node.names]
        elif isinstance(node, ast.Import):
            names = [a.name for a in node.names]
        for n in names:
            if n.split(".")[0] == "beamboot":
                p = n.replace(".", "/") + ".py"
                if os.path.exists(os.path.join(ROOT, p)):
                    module_level_deps(p, seen)


def main():
    seen = set()
    module_level_deps("beamboot/superbeam_gym.py", seen)
    seen.discard("beamboot/__init__.py")
    if os.path.isdir(DEST):
        shutil.rmtree(DEST)
    os.makedirs(os.path.join(DEST, "beamboot"))
    os.makedirs(os.path.join(DEST, "roms"))
    open(os.path.join(DEST, "beamboot", "__init__.py"), "w").write('"""Bundled subset of beamboot for the SUPERBEAM env."""\n')
    for rel in sorted(seen):
        shutil.copy2(os.path.join(ROOT, rel), os.path.join(DEST, rel))
    for r in ROMS:
        shutil.copy2(os.path.join(ROOT, "roms", r), os.path.join(DEST, "roms", r))
    for rel in IN_WEIGHTS:
        os.makedirs(os.path.dirname(os.path.join(DEST, rel)), exist_ok=True)
        shutil.copy2(os.path.join(ROOT, rel), os.path.join(DEST, rel))
    # Lets `pytest tests/` run from vendor/ (as the project root), not inside the env package.
    open(os.path.join(DEST, "pytest.ini"), "w").write("[pytest]\npythonpath = .\n")
    for dirpath, _, files in os.walk(DEST):
        for f in files:
            print(os.path.relpath(os.path.join(dirpath, f), ROOT))


if __name__ == "__main__":
    main()