#
# Proot-Distro - manage proot containers.
#
# Created by Sylirre <sylirre@termux.dev> for Termux project.
# Development assisted by Claude Code (https://claude.ai/code).
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.
#

# Architecture: The stateful interpreter at the core of `proot-distro
# build`. Walks the parsed instructions, dispatches each one to a
# handler (metadata-only handlers live in handlers.py, the RUN
# handler in run_step.py, COPY/ADD in copy_step.py), tracks the
# evolving image config + per-stage state in Stage instances, and
# appends an OCI history entry per dispatched instruction so the
# image_config["history"] array stays in sync with rootfs.diff_ids.

import json
import os
import re

from proot_distro import dirfd
from proot_distro.execenv import is_host_exec_var
from proot_distro.arch import get_device_cpu_arch
from proot_distro.message import C, log_info
from proot_distro.helpers.docker import (
    ARCH_TO_DOCKER, apply_layer, open_required_layer, pull_image,
)
from proot_distro.helpers.dockerfile import expand_vars
from proot_distro.helpers.rootfs import write_hosts, write_resolv_conf
from proot_distro.helpers.build_engine.constants import (
    EXPANDS_VARS, PREDEFINED_ARGS,
)
from proot_distro.helpers.build_engine.dockerignore import load_dockerignore
from proot_distro.helpers.build_engine.errors import BuildError
from proot_distro.helpers.build_engine.handlers import HANDLERS, do_onbuild
from proot_distro.helpers.build_engine.parsing import split_arg
from proot_distro.helpers.build_engine.stage import Stage


_FROM_RE = re.compile(r"^\s*(\S+)(?:\s+AS\s+(\S+))?\s*$", re.IGNORECASE)
_FROM_AS_RE = re.compile(r"\s+AS\s+(\S+)\s*$", re.IGNORECASE)



def _malformed(image_ref: str, what: str):
    """The one refusal shape for a base image config that will not do."""
    raise BuildError(
        f"FROM {image_ref}: the image config is malformed: {what}."
    )


def _cfg_object(value, image_ref: str, what: str) -> dict:
    if not isinstance(value, dict):
        _malformed(image_ref, f"{what} is not an object")
    return value


def _cfg_str_list(value, image_ref: str, what: str) -> list:
    if not isinstance(value, list):
        _malformed(image_ref, f"{what} is not a list")
    for entry in value:
        if not isinstance(entry, str):
            _malformed(image_ref, f"{what} holds an entry that is not "
                                  f"a string")
    return value


def _adopt_image_config(image_config, image_ref: str) -> dict:
    """Return a pulled image's config, held to the shapes read out of it.

    Two things happen here, both at the point a stranger's config is
    *adopted* rather than where it is used: nothing downstream can
    re-seed a value, because a stage's env is seeded from this at FROM,
    `FROM <earlier stage>` deep-copies the whole config, and what the
    build finally stores as its own image's config comes from here too.
    One pass, and every one of those is clean by construction.

    The first is the LD_*/PROOT_* filter over `Env`. An ENV line in the
    user's own Dockerfile still reaches the image config -- it goes
    through do_env, which writes this same list afterwards, and what a
    Dockerfile says about the image it produces is its author's
    business. What no such name reaches from any source is the
    environment proot itself is exec'd with, which run_step refuses
    separately (see proot_distro.execenv for the one rule both apply).

    The second is the shape. Every field below is read back by this
    module or by a handler -- `User` and `Shell` decide what a RUN step
    runs and who as, `WorkingDir` becomes proot's --cwd, `OnBuild` is
    parsed as Dockerfile lines, `Env` seeds the stage, `Cmd` and
    `Entrypoint` are what `run` later executes, `Labels`, `ExposedPorts`
    and `Volumes` are merged into by their handlers, and `history` is
    appended to once per instruction. All of it is a registry's JSON (or
    a manifest-cache entry's, which on Termux sits under the bound
    $TERMUX_PREFIX and is a guest's to compose), and every consumer
    subscripted it as the type OCI says it is: `"config": "x"` was an
    AttributeError at FROM, `"Labels": ["a"]` a ValueError in do_label,
    `"OnBuild": 5` a TypeError, `"history": null` an AttributeError at
    the first instruction -- and `build` catches none of those, so each
    was a traceback.

    A field of the wrong type is a refusal naming it, not a value
    dropped quietly: two of them decide what runs, and the rest would
    otherwise be carried into the image this build publishes. Absent and
    JSON `null` are simply "not set", which is how registries spell an
    empty field; a `null` is removed rather than left standing, since
    `.get(key) or default` and `setdefault(key, default)` do not answer
    alike for one.
    """
    doc = _cfg_object(image_config, image_ref, "the document")

    if doc.get("history") is None:
        doc.pop("history", None)
    elif not isinstance(doc["history"], list):
        _malformed(image_ref, "'history' is not a list")

    if doc.get("rootfs") is None:
        doc.pop("rootfs", None)
    else:
        rootfs = _cfg_object(doc["rootfs"], image_ref, "'rootfs'")
        if rootfs.get("diff_ids") is None:
            rootfs.pop("diff_ids", None)
        else:
            _cfg_str_list(rootfs["diff_ids"], image_ref, "'rootfs.diff_ids'")

    if doc.get("config") is None:
        doc["config"] = {}
        return doc
    cfg = _cfg_object(doc["config"], image_ref, "'config'")

    for key in ("Cmd", "Entrypoint", "OnBuild", "Shell"):
        if cfg.get(key) is None:
            cfg.pop(key, None)
        else:
            _cfg_str_list(cfg[key], image_ref, f"'config.{key}'")

    for key in ("User", "WorkingDir"):
        if cfg.get(key) is None:
            cfg.pop(key, None)
        elif not isinstance(cfg[key], str):
            _malformed(image_ref, f"'config.{key}' is not a string")

    for key in ("ExposedPorts", "Volumes"):
        if cfg.get(key) is None:
            cfg.pop(key, None)
        else:
            # The key set is the whole of what either field says: the
            # spec's value is an empty object, and both handlers write
            # one. Rewritten rather than checked, so an inherited value
            # of any other shape cannot reach the built image either.
            names = _cfg_object(cfg[key], image_ref, f"'config.{key}'")
            cfg[key] = {name: {} for name in names}

    if cfg.get("Labels") is None:
        cfg.pop("Labels", None)
    else:
        labels = _cfg_object(cfg["Labels"], image_ref, "'config.Labels'")
        for value in labels.values():
            if value is not None and not isinstance(value, str):
                _malformed(image_ref, "'config.Labels' holds a value that "
                                      "is not a string")
        # A null label value is the empty string, as it is to every
        # other reader of a map of strings.
        cfg["Labels"] = {k: (v if v is not None else "")
                         for k, v in labels.items()}

    if cfg.get("Env") is None:
        cfg.pop("Env", None)
    else:
        env_list = _cfg_str_list(cfg["Env"], image_ref, "'config.Env'")
        cfg["Env"] = [
            e for e in env_list
            if not ("=" in e and is_host_exec_var(e.partition("=")[0]))
        ]

    return doc


def _stage_rootfs_fd(stage):
    """A private descriptor on *stage*'s rootfs. Raises OSError.

    From the stage's own pinned descriptor when it has one, so the
    caller's walk starts at the inode the engine created rather than at
    `<scratch>/stage-N/rootfs`, a name a previous step's leftovers can
    have re-pointed. The caller owns and closes what comes back.
    """
    if stage.rootfs_fd is not None:
        return dirfd.reopen(stage.rootfs_fd)
    return dirfd.opendir(stage.rootfs_dir)


class BuildEngine:
    """Walks a parsed Dockerfile and produces an OCI image in-place.

    The engine owns the cross-stage state (the global ARG scope, the
    map of named stages so COPY --from= can resolve them, the current
    Stage). Each instruction is dispatched to a handler module
    (handlers.py / copy_step.py / run_step.py) that mutates either
    the engine, the current stage, or the rootfs on disk.
    """

    def __init__(self,
                 build_dir,
                 tmp_root,
                 target_arch_pd,
                 user_build_args,
                 target_stage,
                 verbose,
                 quiet,
                 no_cache,
                 emulator,
                 tmp_root_fd=None):
        self.build_dir = os.path.abspath(build_dir)
        self.tmp_root = tmp_root
        # The scratch root as a descriptor. Every directory the build
        # makes inside it -- a stage's, its rootfs -- is created off this
        # and kept as a descriptor of its own, because the path form is
        # a name anything running as the invoking user can re-point, and
        # a RUN step's leftovers are exactly that. None is the shape a
        # test hands in, and each consumer then keeps the path form.
        self.tmp_root_fd = tmp_root_fd
        self.target_arch_pd = target_arch_pd
        self.target_arch_docker = ARCH_TO_DOCKER.get(
            target_arch_pd, (target_arch_pd, "")
        )[0]
        host_arch = get_device_cpu_arch()
        self.host_arch_docker = ARCH_TO_DOCKER.get(
            host_arch, (host_arch, "")
        )[0]
        self.user_build_args = dict(user_build_args)
        self.target_stage = target_stage
        self.verbose = verbose
        self.quiet = quiet
        self.no_cache = no_cache
        self.emulator = emulator
        self.stages = {}        # name -> Stage
        self.stages_by_idx = []
        self.current = None
        self.global_args = {}   # ARG declared before first FROM
        self.declared_global = set()
        self.ignore_patterns = load_dockerignore(self.build_dir)
        self._stop_after = False
        self._step_no = 0
        self._step_total = 0
        # True only while the base image's ONBUILD triggers are running,
        # so do_env can tell a stranger's ENV from the author's.
        self._firing_onbuild = False
        # LD_*/PROOT_* names the Dockerfile set and the host-side proot
        # exec was refused (run_step._refuse_host_exec). Kept per build
        # so a name is named once rather than once per RUN step.
        self.warned_host_exec = set()

    def _make_stage_dirs(self, idx):
        """Create `stage-<idx>/rootfs` and return both descriptors.

        Made off the scratch root's descriptor, one level at a time with
        O_NOFOLLOW, so the two names cannot be resolved through anything
        a previous stage left behind -- and what the stage keeps is the
        descriptors, not the names it made them under.

        Without a pinned scratch root (a test) this falls back to the
        path form, and the stage carries no descriptors: every consumer
        then behaves exactly as it did before.
        """
        if self.tmp_root_fd is None:
            rootfs_dir = os.path.join(self.tmp_root, f"stage-{idx}", "rootfs")
            os.makedirs(rootfs_dir, exist_ok=True)
            return None, None

        stage_fd = None
        try:
            stage_fd = dirfd.descend_at(
                self.tmp_root_fd, (f"stage-{idx}",), create=True,
            )
            rootfs_fd = dirfd.descend_at(stage_fd, ("rootfs",), create=True)
        except OSError as exc:
            if stage_fd is not None:
                os.close(stage_fd)
            raise BuildError(
                f"cannot create the scratch tree for stage {idx}: {exc}"
            ) from exc
        return stage_fd, rootfs_fd

    def close(self):
        """Release every stage's descriptors. Idempotent."""
        for stage in self.stages_by_idx:
            stage.close()

    # ----- public entry point ----------------------------------------------

    def run(self, instructions):
        """Walk the instruction list and return the chosen stage."""
        self._prescan(instructions)

        self._step_total = len(instructions)
        for instr in instructions:
            self._step_no += 1
            self._announce(instr)
            self._dispatch(instr)
            if self._stop_after:
                break

        if self.current is None:
            raise BuildError("Dockerfile contains no FROM instruction.")

        return self._target()

    # ----- pre-scan --------------------------------------------------------

    def _prescan(self, instructions):
        seen_from = False
        named_stages = []
        for instr in instructions:
            if instr["name"] == "FROM":
                seen_from = True
                value = self._expand_for_from(instr["value"], pre_scan=True)
                m = _FROM_AS_RE.search(value)
                if m:
                    named_stages.append(m.group(1))
            elif instr["name"] == "ARG" and not seen_from:
                key, default = split_arg(instr["value"])
                if not key:
                    continue
                self.declared_global.add(key)
                value = self.user_build_args.get(key, default or "")
                self.global_args[key] = value

        if self.target_stage and self.target_stage not in named_stages:
            raise BuildError(
                f"--target stage '{self.target_stage}' is not defined in "
                f"the Dockerfile (known stages: "
                f"{', '.join(named_stages) or 'none'})."
            )

    def _expand_for_from(self, value, pre_scan=False):
        """Expand variables for a FROM line using global ARGs only."""
        scope = {}
        if pre_scan:
            # During pre-scan we only have CLI-supplied build args for
            # global keys declared so far. Pre-scan happens before the
            # user-arg merge; resolve from the raw map.
            scope.update(self.user_build_args)
        else:
            scope.update(self.global_args)
        self._set_arch_defaults(scope)
        return expand_vars(value, scope)

    def _set_arch_defaults(self, scope: dict) -> None:
        """Populate the TARGET*/BUILD* arch annotations (Docker defaults)."""
        scope.setdefault("TARGETARCH", self.target_arch_docker)
        scope.setdefault("BUILDARCH", self.host_arch_docker)
        scope.setdefault("TARGETOS", "linux")
        scope.setdefault("BUILDOS", "linux")
        scope.setdefault("TARGETPLATFORM", f"linux/{self.target_arch_docker}")
        scope.setdefault("BUILDPLATFORM", f"linux/{self.host_arch_docker}")

    # ----- step banner -----------------------------------------------------

    def _announce(self, instr):
        if self.quiet:
            return
        raw = instr.get("raw", "")
        if len(raw) > 120:
            raw = raw[:117] + "..."
        parts = raw.split(None, 1)
        if len(parts) == 2:
            rendered = f"{C['YELLOW']}{parts[0]}{C['RST']} {parts[1]}"
        elif parts:
            rendered = f"{C['YELLOW']}{parts[0]}{C['RST']}"
        else:
            rendered = ""
        log_info(f"Step {self._step_no}/{self._step_total}: "
                 f"{C['RST']}{rendered}")

    def log(self, text):
        """Emit *text* via log_info() unless `--quiet` is in effect."""
        if self.quiet:
            return
        log_info(text)

    # ----- dispatcher + history -------------------------------------------

    def _dispatch(self, instr):
        name = instr["name"]
        if name == "FROM":
            self._do_from(instr)
            return
        if self.current is None:
            if name == "ARG":
                # Global ARG: already captured by _prescan; no-op here.
                return
            raise BuildError(
                f"Instruction '{name}' before any FROM at line "
                f"{instr['lineno']}."
            )
        if name == "ONBUILD":
            do_onbuild(self, instr)
            self._record_history(instr, layer_added=False)
            return
        if name in EXPANDS_VARS and not instr["exec_form"]:
            instr = self._expand_instruction(instr)

        handler = HANDLERS.get(name)
        if handler is None:
            raise BuildError(
                f"Unsupported instruction '{name}' at line {instr['lineno']}."
            )
        self._run_with_history(handler, instr)

    def _run_with_history(self, handler, instr):
        """Run handler, then append a history entry for the instruction.

        Whether the entry is marked empty_layer depends on whether the
        handler grew `stage.layers`. The OCI image-config spec requires
        the count of non-empty-layer history entries to equal
        len(rootfs.diff_ids); registries (notably Docker Hub) render
        this array in their "Image Layers" UI, so an out-of-date
        history makes built layers invisible.
        """
        layers_before = len(self.current.layers)
        handler(self, instr)
        layer_added = len(self.current.layers) > layers_before
        self._record_history(instr, layer_added=layer_added)

    def _record_history(self, instr, layer_added):
        """Append one entry to image_config["history"] for `instr`.

        `created_by` is the raw Dockerfile line (what Docker Hub
        displays under each step); `created` is fixed to the epoch so
        the image config is reproducible across builds. Entries that
        did not produce a filesystem layer carry empty_layer=true, so
        the count of non-empty entries always equals len(diff_ids).
        """
        entry = {
            "created": "1970-01-01T00:00:00Z",
            "created_by": instr.get("raw") or instr["name"],
        }
        if not layer_added:
            entry["empty_layer"] = True
        cfg = self.current.image_config
        cfg.setdefault("history", []).append(entry)

    def _expand_instruction(self, instr):
        """Return a copy of instr with its `value` variable-expanded."""
        env = self.expansion_scope()
        new = dict(instr)
        value = instr.get("value", "")
        if isinstance(value, str):
            new["value"] = expand_vars(value, env)
        new["flags"] = {
            k: expand_vars(v, env) if isinstance(v, str) else v
            for k, v in instr.get("flags", {}).items()
        }
        return new

    def expansion_scope(self):
        """Variable scope for `${VAR}` expansion inside the current stage.

        Composed in increasing precedence: PREDEFINED_ARGS from the
        host env, the build-time arch annotations (TARGETARCH /
        BUILDARCH / TARGET/BUILDPLATFORM / TARGET/BUILDOS), declared
        ARGs in this stage, and finally ENVs (which win over ARGs by
        Docker semantics).
        """
        scope = {}
        for k in PREDEFINED_ARGS:
            v = os.environ.get(k, "")
            if v:
                scope[k] = v
        self._set_arch_defaults(scope)
        # ARG values declared in this stage so far.
        for k, v in self.current.args.items():
            scope[k] = v
        # ENV always wins over ARG.
        for k, v in self.current.env.items():
            scope[k] = v
        return scope

    # ----- FROM ------------------------------------------------------------

    def _do_from(self, instr):
        # If --target was set and the current stage is the target, the
        # next FROM marks the end of the build for this invocation.
        if (
            self.target_stage
            and self.current is not None
            and self.current.name == self.target_stage
        ):
            self._stop_after = True
            return

        value = self._expand_for_from(
            instr["value"] if isinstance(instr["value"], str) else "",
        )
        m = _FROM_RE.match(value)
        if not m:
            raise BuildError(
                f"Invalid FROM at line {instr['lineno']}: {value!r}"
            )
        base_ref = m.group(1)
        stage_name = m.group(2)

        idx = len(self.stages_by_idx)
        stage_dir = os.path.join(self.tmp_root, f"stage-{idx}")
        rootfs_dir = os.path.join(stage_dir, "rootfs")
        stage_fd, rootfs_fd = self._make_stage_dirs(idx)
        stage = Stage(
            index=idx, name=stage_name, rootfs_dir=rootfs_dir,
            target_arch_pd=self.target_arch_pd,
            dir_fd=stage_fd, rootfs_fd=rootfs_fd,
        )

        # The stage is not registered until the bottom of this method, so
        # a base image that will not resolve would otherwise take its two
        # descriptors with it, past the close() that walks the registered
        # ones.
        try:
            if base_ref.lower() == "scratch":
                stage.image_config = {"config": {}}
            elif base_ref in self.stages:
                self._inherit_from_stage(stage, self.stages[base_ref])
            else:
                self._pull_base_image(stage, base_ref)
        except BaseException:
            stage.close()
            raise

        cfg = stage.image_config.get("config") or {}
        env_list = cfg.get("Env") or []
        for entry in env_list:
            if isinstance(entry, str) and "=" in entry:
                k, _, v = entry.partition("=")
                stage.env[k] = v
        if cfg.get("WorkingDir"):
            stage.workdir = cfg["WorkingDir"] or "/"
        if cfg.get("User"):
            stage.user = cfg["User"]
        if cfg.get("Shell"):
            stage.shell = list(cfg["Shell"])

        # Re-declare global ARGs as available (no value unless declared
        # again in this stage). Docker semantics: they start unset and
        # become visible only after a bare `ARG NAME` re-declares them.
        stage.args = {}
        stage.declared_args = set()

        self.stages_by_idx.append(stage)
        if stage_name:
            self.stages[stage_name] = stage
        # Implicit "by index" lookup for COPY --from=0, --from=1, etc.
        self.stages[str(idx)] = stage
        self.current = stage

        # Configure DNS so RUN apt-get etc. work. Both writers descend
        # from the stage's own descriptor; the isdir() probe is gone
        # because _open_etc already answers "no etc here" by refusing to
        # open it, and asking by name was one more resolution of a name
        # this no longer trusts.
        try:
            write_resolv_conf(rootfs_dir, root_fd=stage.rootfs_fd)
            write_hosts(rootfs_dir, root_fd=stage.rootfs_fd)
        except OSError:
            pass

        # Fire ONBUILD triggers from the base image's config.
        base_onbuild = (
            (stage.image_config.get("config") or {}).get("OnBuild") or []
        )
        if base_onbuild:
            from proot_distro.helpers.dockerfile import parse_dockerfile
            # A trigger that *fires* is always the base image's, never
            # this Dockerfile's: `ONBUILD` here only records one for
            # whoever builds FROM the result. So an ENV among them is a
            # stranger's line, and do_env holds it to the same rule the
            # image's own Env is held to.
            self._firing_onbuild = True
            try:
                for trig in base_onbuild:
                    _, trig_instrs = parse_dockerfile(trig + "\n")
                    for ti in trig_instrs:
                        self._step_no += 1
                        self._announce(ti)
                        if ti["name"] in EXPANDS_VARS and not ti["exec_form"]:
                            ti = self._expand_instruction(ti)
                        h = HANDLERS.get(ti["name"])
                        if h is None:
                            raise BuildError(
                                f"ONBUILD trigger uses unsupported "
                                f"instruction '{ti['name']}'."
                            )
                        self._run_with_history(h, ti)
            finally:
                self._firing_onbuild = False

    def _inherit_from_stage(self, new_stage, parent):
        """Apply parent's layers to new_stage.rootfs_dir; copy config.

        The deep-copy via JSON round-trip carries the parent's
        `history` array along with the rest of image_config, so the
        new stage starts with the inherited entries and subsequent
        instructions append to the same list.
        """
        new_stage.image_config = json.loads(
            json.dumps(parent.image_config)
        )
        # The stage's own descriptor, not its name: the directory was
        # made by this process, but by the time a second FROM re-applies
        # layers into it a previous stage's RUN steps have run, and what
        # they leave running is the invoking user.
        rootfs_fd = _stage_rootfs_fd(new_stage)
        try:
            for layer in parent.layers:
                # Verified, not merely present: these layers were produced
                # by this build, so a blob that no longer hashes to its
                # digest cannot be refetched and must not be applied.
                try:
                    layer_fd = open_required_layer(
                        layer["digest"],
                        what=(f"Layer of stage "
                              f"'{parent.name or parent.index}':"),
                    )
                except RuntimeError as exc:
                    raise BuildError(str(exc)) from exc
                try:
                    apply_layer(layer_fd, rootfs_fd,
                                digest=layer["digest"])
                finally:
                    os.close(layer_fd)
        finally:
            os.close(rootfs_fd)
        new_stage.layers = list(parent.layers)
        new_stage.parent_layer_digest = parent.parent_layer_digest

    def _pull_base_image(self, stage, image_ref):
        """Use helpers.docker.pull_image to populate the stage rootfs."""
        log_info(f"Pulling base image '{image_ref}' ({self.target_arch_pd})...")
        try:
            rootfs_fd = _stage_rootfs_fd(stage)
        except OSError as exc:
            raise BuildError(f"FROM {image_ref}: {exc}") from exc
        try:
            meta = pull_image(image_ref, rootfs_fd, self.target_arch_pd)
        except RuntimeError as exc:
            raise BuildError(f"FROM {image_ref}: {exc}") from exc
        finally:
            os.close(rootfs_fd)

        stage.image_config = _adopt_image_config(
            meta.get("image_config") or {"config": {}}, image_ref,
        )
        manifest = meta.get("manifest") or {}
        config_diff_ids = (
            (stage.image_config.get("rootfs") or {}).get("diff_ids") or []
        )
        stage.layers = []
        for i, layer in enumerate(manifest.get("layers", [])):
            # pull_image has already held the manifest to manifest_layers'
            # rule, so `layers` is a list of descriptors and each digest
            # is a string. A `size` is not part of that promise, and this
            # one is copied into the manifest the build publishes: read a
            # non-int as absent, the way every other reader of the field
            # does.
            digest = layer.get("digest", "")
            size = layer.get("size", 0)
            if not isinstance(size, int) or isinstance(size, bool):
                size = 0
            diff_id = (
                config_diff_ids[i]
                if i < len(config_diff_ids) else digest
            )
            stage.layers.append(
                {"digest": digest, "size": size, "diff_id": diff_id}
            )
        if stage.layers:
            stage.parent_layer_digest = stage.layers[-1]["digest"]

    # ----- target stage resolution ----------------------------------------

    def _target(self):
        if self.target_stage:
            stage = self.stages.get(self.target_stage)
            if stage is None:
                raise BuildError(
                    f"--target stage '{self.target_stage}' was not built."
                )
            return stage
        return self.current
