compose_module_layers.py

#!/usr/bin/env python3
# Copyright (c) 2026 Carnegie Mellon University
# SPDX-License-Identifier: BSD-3-Clause-Clear
"""Compose per-module Docker layers into per-host image plans (RFC #379 §6, Phase P4).

Trunk publishes one signed base image per host type per version; modules bring
their dependencies in three tiers declared in ``module.yaml``:

- **tier 1** — ``deps.apt`` / ``deps.pip`` lists → one generated Dockerfile
  layer per module (``RUN apt-get install …`` / ``RUN pip3 install …``), so a
  module's dep change invalidates only its own layer cache.
- **tier 2** — ``dockerfile:`` → the module's ``Dockerfile.module``, written
  against ``ARG BASE_IMAGE``, built with BASE_IMAGE = the previous chain link.
- **tier 3** — ``overlay_image:`` → a prebuilt overlay. Used **as-is** only
  when it is the sole docker-relevant module for that host; in every other
  composition the module must also carry a ``dockerfile:`` (fragment = source
  of truth, overlay = cache) — RFC #379 §6, failure mode 2.

The chain is deterministic: modules sorted by name within each tier, tiers in
order 1 → 2 → 3-as-build, grouped per target host (``robot`` | ``gcs`` |
``isaac-sim`` | ``ms-airsim``).

**Zero-module identity rule (hard requirement):** when no module contributes a
docker-relevant declaration, every host's plan is exactly
``{base_image: <trunk tag>, steps: [], final_tag: <trunk tag>}`` — no
``Dockerfile.composed`` is generated and no ``image:`` override is ever added
to the generated compose file. A module-free (or dep-free) checkout uses
today's images, byte-identically.

Outputs (default mode — plan + lock, no docker calls):

- ``.airstack/generated/layer_plan.json`` —
  ``{host: {base_image, steps: [{module, tier, dockerfile|null, dep_hash}], final_tag}}``
- ``.airstack/generated/layers/<host>/Dockerfile.composed`` — the tier-1 stage,
  only for hosts that have tier-1 steps
- ``modules.lock`` at the repo root (gitignored) — per module
  ``{name, pin, dep_hash, targets}`` plus ``plan_hash``; serialization is
  deterministic (sorted keys/names), so identical inputs → byte-identical lock

``final_tag`` naming extends the trunk scheme:
``${PROJECT_DOCKER_REGISTRY}/${PROJECT_NAME}:v${VERSION}_<host-suffix>-m<plan_hash[:8]>``.
Composed tags are **per-checkout artifacts** — they are never pushed by trunk
CI and never enter the docker-build.yml publish pipeline.

CLI::

    compose_module_layers.py [--project-root DIR]      # plan + modules.lock (default)
    compose_module_layers.py --check-conflicts [...]   # static apt/pip conflict gate
    compose_module_layers.py --build [...]             # run the docker build chain
                                                       # + compose image overrides
                                                       # (CI/orchestrator only)

``--check-conflicts`` is doctor hard gate #1 (RFC #379 §4): duplicate apt/pip
packages with *different* version specs across modules on the same host fail,
naming the fighting modules. Same-spec duplicates and unpinned duplicates pass.

Exit 0 on success; 1 on a plan error or a dependency conflict.
"""
import argparse
import hashlib
import json
import os
import re
import shutil
import subprocess
import sys
from pathlib import Path

import yaml

MANIFEST_NAME = "module.yaml"
MODULES_REL = Path("modules")
REPOS_FILE_NAME = "modules.repos"
LOCK_REL = Path("modules.lock")
PLAN_REL = Path(".airstack/generated/layer_plan.json")
LAYERS_REL = Path(".airstack/generated/layers")
GENERATED_COMPOSE_REL = Path(".airstack/generated/docker-compose.modules.yaml")

# Target hosts a module may declare (module.schema.json `targets` enum) mapped
# to the trunk image-tag suffix scheme (see robot/docker/docker-compose.yaml,
# gcs/docker/gcs-base-docker-compose.yaml, simulation/*/docker/docker-compose.yaml).
# The robot suffix carries DOCKER_IMAGE_BUILD_MODE, the others do not — that is
# the existing trunk scheme, mirrored here on purpose.
HOSTS = ("robot", "gcs", "isaac-sim", "ms-airsim")

# Compose services that get an `image:` override under --build, per host.
# robot-l4t is deliberately absent: its image chains from robot-l4t-stack-base
# (aarch64), so pointing it at an x86-64 composed image would be wrong. Override
# with AIRSTACK_MODULE_LAYER_ROBOT_SERVICES if your checkout differs.
HOST_SERVICES = {
    "robot": ("robot-desktop",),
    "gcs": ("gcs",),
    "isaac-sim": ("isaac-sim",),
    "ms-airsim": ("ms-airsim",),
}

COMPOSED_HEADER = """\
# GENERATED by tools/compose_module_layers.py (`airstack module sync`) — DO NOT EDIT.
# Tier-1 module dependency layers (RFC #379 §6): one RUN per module per package
# manager, so a module's dep change invalidates only its own layer cache.
# Built with: docker build --build-arg BASE_IMAGE=<previous chain link> …
"""

COMPOSE_HEADER = """\
# GENERATED by tools/module_overlay.py + tools/compose_module_layers.py — DO NOT EDIT.
# `image:` keys below were added by `compose_module_layers.py --build` and point
# services at the composed module-layer images for this checkout. Re-running
# `airstack module sync` regenerates the file without them (plan-only).
"""


def log(msg):
    print(f"[layers] {msg}")


class LayerPlanError(Exception):
    pass


# ── inputs: .env, modules.repos pins, module manifests ───────────────────────

def read_env(root):
    """Parse the top-level .env (KEY=VALUE lines, quotes stripped, comments ignored)."""
    env = {}
    env_path = root / ".env"
    if not env_path.is_file():
        return env
    for line in env_path.read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        key, _, raw = line.partition("=")
        raw = raw.strip()
        if raw[:1] in ("'", '"'):
            quote = raw[0]
            end = raw.find(quote, 1)
            raw = raw[1:end] if end > 0 else raw[1:]
        else:
            raw = raw.split("#", 1)[0].strip()
        env[key.strip()] = raw
    return env


def read_pins(root):
    """{module_name: pin} from modules.repos ('local' for x-local-modules entries)."""
    pins = {}
    repos_path = root / REPOS_FILE_NAME
    if not repos_path.is_file():
        return pins
    data = yaml.safe_load(repos_path.read_text(encoding="utf-8")) or {}
    for name, repo in (data.get("repositories") or {}).items():
        pins[name] = str((repo or {}).get("version", "?"))
    for entry in data.get("x-local-modules") or []:
        pins[entry.get("name")] = "local"
    return pins


def discover_modules(root):
    """Map <checkout-dir-name> → manifest for every modules/<name>/module.yaml."""
    modules = {}
    modules_dir = root / MODULES_REL
    if not modules_dir.is_dir():
        return modules
    for child in sorted(modules_dir.iterdir()):
        if not child.is_dir():  # follows symlinks; dangling links land here
            continue
        manifest_path = child / MANIFEST_NAME
        if not manifest_path.is_file():
            continue
        try:
            manifest = yaml.safe_load(manifest_path.read_text(encoding="utf-8"))
        except yaml.YAMLError as exc:
            raise LayerPlanError(f"{manifest_path}: invalid YAML: {exc}")
        if not isinstance(manifest, dict):
            raise LayerPlanError(f"{manifest_path}: manifest top level must be a mapping")
        modules[child.name] = manifest
    return modules


# ── per-module docker declarations ───────────────────────────────────────────

def module_decl(root, name, manifest):
    """Normalize one module's docker-relevant declarations."""
    deps = manifest.get("deps") or {}
    decl = {
        "apt": [str(p) for p in (deps.get("apt") or [])],
        "pip": [str(p) for p in (deps.get("pip") or [])],
        "dockerfile": manifest.get("dockerfile") or None,
        "overlay_image": manifest.get("overlay_image") or None,
        "targets": [t for t in (manifest.get("targets") or []) if t in HOSTS],
    }
    if decl["dockerfile"]:
        dockerfile_path = root / MODULES_REL / name / decl["dockerfile"]
        if not dockerfile_path.is_file():
            raise LayerPlanError(
                f"{name}: declared dockerfile not found: {dockerfile_path}"
            )
        decl["dockerfile_bytes"] = dockerfile_path.read_bytes()
        if b"BASE_IMAGE" not in decl["dockerfile_bytes"]:
            log(f"warning: {name}/{decl['dockerfile']} does not reference BASE_IMAGE "
                "— tier-2 fragments must build against ARG BASE_IMAGE, never a fixed base")
    else:
        decl["dockerfile_bytes"] = None
    return decl


def is_docker_relevant(decl):
    return bool(decl["apt"] or decl["pip"] or decl["dockerfile"] or decl["overlay_image"])


def dep_hash(decl):
    """sha256 over canonicalized deps + dockerfile bytes + overlay_image."""
    payload = {
        "apt": sorted(decl["apt"]),
        "pip": sorted(decl["pip"]),
        "dockerfile": decl["dockerfile"],
        "dockerfile_sha256": (
            hashlib.sha256(decl["dockerfile_bytes"]).hexdigest()
            if decl["dockerfile_bytes"] is not None else None
        ),
        "overlay_image": decl["overlay_image"],
    }
    canonical = json.dumps(payload, sort_keys=True, separators=(",", ":"))
    return hashlib.sha256(canonical.encode("utf-8")).hexdigest()


# ── conflict gate (doctor hard gate #1 — RFC #379 §4/§6) ─────────────────────

_PKG_NAME_RE = re.compile(r"^([A-Za-z0-9][A-Za-z0-9._+-]*)\s*(.*)$")


def _split_spec(entry, manager):
    """'tabulate==0.9.0' → ('tabulate', '==0.9.0'); apt 'pkg=1.2' → ('pkg', '=1.2')."""
    entry = entry.strip()
    if manager == "apt":
        name, _, spec = entry.partition("=")
        return name.strip(), ("=" + spec.strip()) if spec else ""
    match = _PKG_NAME_RE.match(entry)
    if not match:
        return entry, ""
    # PEP 503 name normalization for pip
    name = re.sub(r"[-_.]+", "-", match.group(1)).lower()
    return name, match.group(2).strip()


def find_conflicts(decls):
    """Static text analysis: same package, different non-empty specs, same host.

    Returns a list of human-readable conflict strings (empty = clean). Unpinned
    duplicates and identical-spec duplicates are fine — only *different* pins
    on the same package fight.
    """
    conflicts = []
    for host in HOSTS:
        for manager in ("apt", "pip"):
            by_name = {}  # pkg → {spec: [module, …]}
            for module in sorted(decls):
                decl = decls[module]
                if host not in decl["targets"]:
                    continue
                for entry in decl[manager]:
                    name, spec = _split_spec(entry, manager)
                    by_name.setdefault(name, {}).setdefault(spec, []).append(
                        (module, entry)
                    )
            for name, specs in sorted(by_name.items()):
                pinned = {s: mods for s, mods in specs.items() if s}
                if len(pinned) <= 1:
                    continue
                sides = "; ".join(
                    f"'{entry}' ({module})"
                    for spec in sorted(pinned)
                    for module, entry in pinned[spec]
                )
                conflicts.append(
                    f"{manager} package '{name}' pinned differently for host "
                    f"'{host}': {sides}"
                )
    return conflicts


# ── plan construction ────────────────────────────────────────────────────────

def base_image(host, env):
    registry = env.get("PROJECT_DOCKER_REGISTRY", "airstack")
    project = env.get("PROJECT_NAME", "airstack")
    version = env.get("VERSION", "0.0.0")
    return f"{registry}/{project}:v{version}_{host_suffix(host, env)}"


def host_suffix(host, env):
    if host == "robot":
        return f"robot-x86-64_{env.get('DOCKER_IMAGE_BUILD_MODE', 'dev')}"
    return host


def build_plan(root, modules, env):
    """Compute (plan, lock_entries, plan_hash, decls).

    plan: {host: {base_image, steps, final_tag}} for every host — hosts with no
    docker-relevant modules get the identity entry (steps: [], final == base).
    """
    pins = read_pins(root)
    decls = {name: module_decl(root, name, manifest) for name, manifest in modules.items()}

    lock_entries = [
        {
            "name": name,
            "pin": pins.get(name, "local"),
            "dep_hash": dep_hash(decls[name]),
            "targets": sorted(decls[name]["targets"]),
        }
        for name in sorted(modules)
    ]
    plan_hash = hashlib.sha256(
        json.dumps(lock_entries, sort_keys=True, separators=(",", ":")).encode("utf-8")
    ).hexdigest()

    plan = {}
    for host in HOSTS:
        base = base_image(host, env)
        on_host = {n: d for n, d in decls.items() if host in d["targets"]}
        tier1 = sorted(n for n, d in on_host.items() if d["apt"] or d["pip"])
        tier2 = sorted(n for n, d in on_host.items()
                       if d["dockerfile"] and not d["overlay_image"])
        tier3 = sorted(n for n, d in on_host.items() if d["overlay_image"])

        # Tier-3 rule (RFC #379 §6, failure mode 2): a prebuilt overlay was
        # published FROM the plain trunk base, so docker cannot merge it into a
        # locally-built chain. It is used as-is only when it is the sole
        # docker-relevant module for this host; in any other composition its
        # module must also carry a dockerfile (fragment = source of truth,
        # overlay = cache) and the fragment is built in chain order.
        as_is = None
        for name in tier3:
            if on_host[name]["dockerfile"]:
                continue  # tier-3-as-build: fragment is the source of truth
            sole = (
                len(tier3) == 1
                and not tier1
                and not tier2
            )
            if sole:
                as_is = name
            else:
                others = sorted(set(tier1) | set(tier2) | set(tier3) - {name})
                raise LayerPlanError(
                    f"host '{host}': module '{name}' declares overlay_image "
                    f"'{on_host[name]['overlay_image']}' without a dockerfile, but it is "
                    f"not the sole docker-relevant module (also composing: {', '.join(others)}). "
                    "A prebuilt overlay is used as-is only when it stands alone; otherwise "
                    "the module must also carry a Dockerfile.module — fragment = source of "
                    "truth, overlay = cache (RFC #379 §6)."
                )

        steps = []
        for name in tier1:
            steps.append({"module": name, "tier": 1, "dockerfile": None,
                          "dep_hash": dep_hash(on_host[name])})
        for name in tier2:
            steps.append({"module": name, "tier": 2,
                          "dockerfile": str(MODULES_REL / name / on_host[name]["dockerfile"]),
                          "dep_hash": dep_hash(on_host[name])})
        for name in tier3:
            decl = on_host[name]
            steps.append({"module": name, "tier": 3,
                          "dockerfile": (str(MODULES_REL / name / decl["dockerfile"])
                                         if decl["dockerfile"] else None),
                          "dep_hash": dep_hash(decl)})

        if not steps:
            final_tag = base                      # identity: base images unchanged
        elif as_is is not None:
            final_tag = on_host[as_is]["overlay_image"]  # pulled, never rebuilt
        else:
            final_tag = f"{base}-m{plan_hash[:8]}"

        plan[host] = {"base_image": base, "steps": steps, "final_tag": final_tag}

    return plan, lock_entries, plan_hash, decls


def render_composed_dockerfile(host, plan_entry, decls):
    """The tier-1 stage: ARG BASE_IMAGE chain link, one RUN per module per manager."""
    lines = [COMPOSED_HEADER, "ARG BASE_IMAGE", "FROM ${BASE_IMAGE}", ""]
    for step in plan_entry["steps"]:
        if step["tier"] != 1:
            continue
        decl = decls[step["module"]]
        lines.append(f"# module: {step['module']} (tier 1, dep_hash {step['dep_hash'][:12]})")
        if decl["apt"]:
            pkgs = " ".join(_shell_quote(p) for p in decl["apt"])
            lines.append(
                "RUN apt-get update && "
                f"apt-get install -y --no-install-recommends {pkgs} && "
                "rm -rf /var/lib/apt/lists/*"
            )
        if decl["pip"]:
            pkgs = " ".join(_shell_quote(p) for p in decl["pip"])
            lines.append(f"RUN pip3 install --no-cache-dir --break-system-packages {pkgs}")
        lines.append("")
    return "\n".join(lines).rstrip("\n") + "\n"


def _shell_quote(value):
    if re.fullmatch(r"[A-Za-z0-9._+=:@/-]+", value):
        return value
    return "'" + value.replace("'", "'\\''") + "'"


# ── output writing ───────────────────────────────────────────────────────────

def _write_if_changed(path, text):
    if path.exists() and path.read_text(encoding="utf-8") == text:
        return False
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(text, encoding="utf-8")
    return True


def write_outputs(root, plan, lock_entries, plan_hash, decls):
    plan_path = root / PLAN_REL
    lock_path = root / LOCK_REL
    layers_dir = root / LAYERS_REL

    if not lock_entries:
        # No modules at all — remove every generated layer artifact.
        removed = False
        for path in (plan_path, lock_path):
            if path.exists():
                path.unlink()
                removed = True
        if layers_dir.is_dir():
            shutil.rmtree(layers_dir)
            removed = True
        for parent in (plan_path.parent,):
            try:
                parent.rmdir()
            except OSError:
                pass
        if removed:
            log("no modules — removed generated layer artifacts")
        return

    _write_if_changed(plan_path, json.dumps(plan, indent=2, sort_keys=True) + "\n")
    lock = {"modules": lock_entries, "plan_hash": plan_hash}
    _write_if_changed(lock_path, json.dumps(lock, indent=2, sort_keys=True) + "\n")

    # Dockerfile.composed only for hosts with tier-1 steps; prune the rest.
    wanted = {}
    for host, entry in plan.items():
        if any(step["tier"] == 1 for step in entry["steps"]):
            wanted[host] = render_composed_dockerfile(host, entry, decls)
    if layers_dir.is_dir():
        for child in layers_dir.iterdir():
            if child.name not in wanted:
                shutil.rmtree(child) if child.is_dir() else child.unlink()
        try:
            layers_dir.rmdir()
        except OSError:
            pass
    for host, text in wanted.items():
        _write_if_changed(layers_dir / host / "Dockerfile.composed", text)


def log_summary(plan, decls):
    relevant = sorted(n for n, d in decls.items() if is_docker_relevant(d))
    if not relevant:
        log("module layers: 0 docker-relevant modules — base images unchanged")
        return
    log(f"module layers: {len(relevant)} docker-relevant module(s): {', '.join(relevant)}")
    for host in HOSTS:
        entry = plan[host]
        if entry["steps"]:
            log(f"  {host}: {len(entry['steps'])} step(s) → {entry['final_tag']}")
        else:
            log(f"  {host}: base image unchanged")


# ── --build: run the chain (CI/orchestrator path — needs docker) ─────────────

def robot_services():
    raw = os.environ.get("AIRSTACK_MODULE_LAYER_ROBOT_SERVICES")
    if raw:
        return tuple(s.strip() for s in raw.split(",") if s.strip())
    return HOST_SERVICES["robot"]


def run_builds(root, plan, decls):
    for host in sorted(plan):
        entry = plan[host]
        if not entry["steps"]:
            continue
        if entry["final_tag"] == _sole_overlay_ref(entry, decls):
            log(f"{host}: pulling prebuilt overlay as-is: {entry['final_tag']}")
            subprocess.run(["docker", "pull", entry["final_tag"]], check=True)
            continue
        builds = []
        composed = root / LAYERS_REL / host / "Dockerfile.composed"
        if any(step["tier"] == 1 for step in entry["steps"]):
            builds.append((composed, composed.parent))
        for step in entry["steps"]:
            if step["dockerfile"]:
                dockerfile = root / step["dockerfile"]
                context = Path(os.path.realpath(root / MODULES_REL / step["module"]))
                builds.append((dockerfile, context))
        current = entry["base_image"]
        for index, (dockerfile, context) in enumerate(builds):
            last = index == len(builds) - 1
            tag = entry["final_tag"] if last else f"{entry['final_tag']}-l{index}"
            cmd = ["docker", "build", "-f", str(dockerfile),
                   "--build-arg", f"BASE_IMAGE={current}", "-t", tag, str(context)]
            log(f"{host}: {' '.join(cmd)}")
            subprocess.run(cmd, check=True)
            current = tag


def _sole_overlay_ref(entry, decls):
    if len(entry["steps"]) != 1:
        return None
    step = entry["steps"][0]
    if step["tier"] == 3 and step["dockerfile"] is None:
        return decls[step["module"]]["overlay_image"]
    return None


def apply_compose_image_overrides(root, plan):
    """Point host services at the composed tags in the generated compose override.

    Only called under --build. The plan-only path never writes `image:` keys —
    that is the zero-module identity guarantee.
    """
    compose_path = root / GENERATED_COMPOSE_REL
    data = {}
    if compose_path.exists():
        data = yaml.safe_load(compose_path.read_text(encoding="utf-8")) or {}
    services = data.setdefault("services", {})
    changed = False
    for host, entry in plan.items():
        if entry["final_tag"] == entry["base_image"]:
            continue
        host_svcs = robot_services() if host == "robot" else HOST_SERVICES[host]
        for service in host_svcs:
            services.setdefault(service, {})["image"] = entry["final_tag"]
            changed = True
    if not changed:
        return
    compose_path.parent.mkdir(parents=True, exist_ok=True)
    compose_path.write_text(
        COMPOSE_HEADER + yaml.safe_dump(data, sort_keys=False, default_flow_style=False),
        encoding="utf-8",
    )
    log(f"wrote image overrides into {compose_path}")


# ── entry point ──────────────────────────────────────────────────────────────

def run(root, mode="plan"):
    root = Path(root).resolve()
    env = read_env(root)
    modules = discover_modules(root)

    if mode == "check-conflicts":
        decls = {name: module_decl(root, name, manifest)
                 for name, manifest in modules.items()}
        conflicts = find_conflicts(decls)
        for conflict in conflicts:
            log(f"CONFLICT: {conflict}")
        if conflicts:
            log(f"{len(conflicts)} dependency conflict(s) — modules cannot compose "
                "into one image (RFC #379 §6). Fix or align the pins above.")
            return 1
        log(f"no apt/pip conflicts across {len(modules)} module(s)")
        return 0

    plan, lock_entries, plan_hash, decls = build_plan(root, modules, env)
    write_outputs(root, plan, lock_entries, plan_hash, decls)
    log_summary(plan, decls)

    if mode == "build":
        run_builds(root, plan, decls)
        apply_compose_image_overrides(root, plan)
    return 0


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--project-root", default=str(Path(__file__).resolve().parents[1]),
                        help="AirStack checkout root (default: parent of tools/)")
    group = parser.add_mutually_exclusive_group()
    group.add_argument("--check-conflicts", action="store_true",
                       help="report apt/pip packages pinned differently across modules; "
                            "exit 1 on conflict (doctor hard gate #1 — sync fails on this)")
    group.add_argument("--build", action="store_true",
                       help="after planning, run the docker build chain and point the "
                            "generated compose override at the composed tags "
                            "(CI/orchestrator path — requires docker)")
    args = parser.parse_args(argv)

    mode = "check-conflicts" if args.check_conflicts else "build" if args.build else "plan"
    try:
        return run(args.project_root, mode=mode)
    except LayerPlanError as exc:
        log(f"ERROR: {exc}")
        return 1
    except subprocess.CalledProcessError as exc:
        log(f"ERROR: build command failed: {exc}")
        return 1


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