zlayer_builder/dockerfile/

parser.rs

//! Dockerfile parser
//!
//! This module provides functionality to parse Dockerfiles into a structured representation
//! using the `dockerfile-parser` crate as the parsing backend.

use std::collections::HashMap;
use std::path::Path;
use std::str::FromStr;

use dockerfile_parser::{Dockerfile as RawDockerfile, Instruction as RawInstruction};
use serde::{Deserialize, Serialize};
use zlayer_types::ImageReference;

use crate::error::{BuildError, Result};

use super::instruction::{
    AddInstruction, ArgInstruction, CopyInstruction, EnvInstruction, ExposeInstruction,
    ExposeProtocol, HealthcheckInstruction, Instruction, RunInstruction, ShellOrExec,
};

/// A Dockerfile `FROM` target.
///
/// `FROM` references can resolve to one of three things in a Dockerfile:
/// an OCI image (the common case), a previous stage in a multi-stage
/// build (e.g. `FROM builder AS final`), or the special `scratch`
/// pseudo-image. This enum captures all three. For non-Dockerfile call
/// sites (image registry lookups, toolchain detection, etc.) use
/// [`zlayer_types::ImageReference`] directly — the bare OCI ref type
/// without the Dockerfile-only variants.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum DockerfileFromTarget {
    /// An OCI image reference (canonical OCI grammar).
    Image(ImageReference),
    /// A reference to another stage in this multi-stage build.
    Stage(String),
    /// The special `scratch` pseudo-image.
    Scratch,
}

impl DockerfileFromTarget {
    /// Parse a raw `FROM` target string.
    ///
    /// Recognizes `scratch` (case-insensitive), then attempts an OCI
    /// reference parse via [`ImageReference::from_str`]. If parsing
    /// succeeds, the result is an [`Self::Image`]; otherwise the
    /// input is treated as a [`Self::Stage`] reference.
    ///
    /// Note that the OCI grammar accepts bare names like `alpine` as
    /// valid image references, so disambiguation between an image
    /// and a multi-stage stage reference must happen post-hoc at the
    /// call site by consulting the set of known stage names.
    #[must_use]
    pub fn parse(s: &str) -> Self {
        let s = s.trim();

        if s.eq_ignore_ascii_case("scratch") {
            return Self::Scratch;
        }

        match ImageReference::from_str(s) {
            Ok(r) => Self::Image(r),
            Err(_) => Self::Stage(s.to_string()),
        }
    }

    /// Returns true if this is a stage reference.
    #[must_use]
    pub fn is_stage(&self) -> bool {
        matches!(self, Self::Stage(_))
    }

    /// Returns true if this is the `scratch` pseudo-image.
    #[must_use]
    pub fn is_scratch(&self) -> bool {
        matches!(self, Self::Scratch)
    }
}

impl std::fmt::Display for DockerfileFromTarget {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Image(r) => write!(f, "{r}"),
            Self::Stage(name) => f.write_str(name),
            Self::Scratch => f.write_str("scratch"),
        }
    }
}

/// A single stage in a multi-stage Dockerfile
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Stage {
    /// Stage index (0-based)
    pub index: usize,

    /// Optional stage name (from `AS name`)
    pub name: Option<String>,

    /// The base image for this stage
    pub base_image: DockerfileFromTarget,

    /// Optional platform specification (e.g., "linux/amd64")
    pub platform: Option<String>,

    /// Instructions in this stage (excluding the FROM)
    pub instructions: Vec<Instruction>,
}

impl Stage {
    /// Returns the stage identifier (name if present, otherwise index as string)
    #[must_use]
    pub fn identifier(&self) -> String {
        self.name.clone().unwrap_or_else(|| self.index.to_string())
    }

    /// Returns true if this stage matches the given name or index
    #[must_use]
    pub fn matches(&self, name_or_index: &str) -> bool {
        if let Some(ref name) = self.name {
            if name == name_or_index {
                return true;
            }
        }

        if let Ok(idx) = name_or_index.parse::<usize>() {
            return idx == self.index;
        }

        false
    }
}

/// Expand Docker build-arg references in a `FROM` target string.
///
/// Supports the forms Docker accepts in `FROM` lines: `${VAR}`,
/// `${VAR:-default}` (default when unset or empty), `${VAR:+alt}` (alt
/// when set and non-empty), and bare `$VAR`. Unknown variables expand to
/// the empty string, matching `docker build`.
fn expand_from_args(input: &str, vars: &HashMap<String, String>) -> String {
    let mut out = String::with_capacity(input.len());
    let mut chars = input.char_indices().peekable();
    while let Some((_, c)) = chars.next() {
        if c != '$' {
            out.push(c);
            continue;
        }
        match chars.peek() {
            Some(&(_, '{')) => {
                chars.next(); // consume '{'
                let mut body = String::new();
                let mut closed = false;
                for (_, bc) in chars.by_ref() {
                    if bc == '}' {
                        closed = true;
                        break;
                    }
                    body.push(bc);
                }
                if !closed {
                    // Unterminated `${...` — emit verbatim.
                    out.push_str("${");
                    out.push_str(&body);
                    continue;
                }
                if let Some((name, default)) = body.split_once(":-") {
                    match vars.get(name).filter(|v| !v.is_empty()) {
                        Some(v) => out.push_str(v),
                        None => out.push_str(default),
                    }
                } else if let Some((name, alt)) = body.split_once(":+") {
                    if vars.get(name).is_some_and(|v| !v.is_empty()) {
                        out.push_str(alt);
                    }
                } else {
                    out.push_str(vars.get(&body).map_or("", String::as_str));
                }
            }
            Some(&(_, next)) if next.is_ascii_alphabetic() || next == '_' => {
                let mut name = String::new();
                while let Some(&(_, nc)) = chars.peek() {
                    if nc.is_ascii_alphanumeric() || nc == '_' {
                        name.push(nc);
                        chars.next();
                    } else {
                        break;
                    }
                }
                out.push_str(vars.get(&name).map_or("", String::as_str));
            }
            // BuildKit escape: `$$` is a literal dollar sign.
            Some(&(_, '$')) => {
                chars.next();
                out.push('$');
            }
            _ => out.push('$'),
        }
    }
    out
}

/// A parsed Dockerfile
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Dockerfile {
    /// Global ARG instructions that appear before the first FROM
    pub global_args: Vec<ArgInstruction>,

    /// Build stages
    pub stages: Vec<Stage>,
}

impl Dockerfile {
    /// Expand pre-FROM `ARG`s in every stage's `FROM` target.
    ///
    /// Docker semantics: only ARGs declared BEFORE the first `FROM`
    /// participate in `FROM`-line expansion (`FROM ${BASE_IMAGE}` /
    /// `FROM img:${TAG:-latest}`), with `--build-arg` values overriding
    /// their defaults. The parser can't do this (it has no build args), so
    /// targets containing `$` end up classified as [`DockerfileFromTarget::Stage`]
    /// — and struct-consuming backends then fail with `Stage '${BASE_IMAGE}'
    /// not found`. Call this with the effective build args before handing
    /// the Dockerfile to a backend.
    ///
    /// Expanded targets are re-classified: `scratch`, a previously declared
    /// stage name, an OCI image reference, or (still) a stage string. A
    /// target that expands to an empty string is left untouched so the
    /// eventual error names the unexpanded variable instead of a blank.
    pub fn resolve_from_args(&mut self, build_args: &HashMap<String, String>) {
        // Effective FROM-scope variables: declared global ARGs only,
        // defaults overridden by matching build args (an undeclared build
        // arg does NOT leak into FROM lines, per Docker).
        let mut vars: HashMap<String, String> = HashMap::new();
        for arg in &self.global_args {
            let value = build_args
                .get(&arg.name)
                .cloned()
                .or_else(|| arg.default.clone())
                .unwrap_or_default();
            vars.insert(arg.name.clone(), value);
        }

        let mut known_stage_names: std::collections::HashSet<String> =
            std::collections::HashSet::new();
        for stage in &mut self.stages {
            if let DockerfileFromTarget::Stage(raw) = &stage.base_image {
                if raw.contains('$') {
                    let expanded = expand_from_args(raw, &vars);
                    if !expanded.trim().is_empty() {
                        let mut target = DockerfileFromTarget::parse(&expanded);
                        // Same post-hoc stage promotion as `from_raw`: a bare
                        // name that matches an earlier stage alias is a stage
                        // reference even though it parses as an OCI ref.
                        if matches!(target, DockerfileFromTarget::Image(_))
                            && known_stage_names.contains(expanded.trim())
                        {
                            target = DockerfileFromTarget::Stage(expanded.trim().to_string());
                        }
                        stage.base_image = target;
                    }
                }
            }
            if let Some(name) = &stage.name {
                known_stage_names.insert(name.clone());
            }
        }
    }

    /// Parse a Dockerfile from a string
    ///
    /// # Errors
    ///
    /// Returns an error if the Dockerfile content is malformed or contains invalid instructions.
    pub fn parse(content: &str) -> Result<Self> {
        let raw = RawDockerfile::parse(content).map_err(|e| BuildError::DockerfileParse {
            message: e.to_string(),
            line: 1,
        })?;

        Self::from_raw(raw)
    }

    /// Parse a Dockerfile from a file
    ///
    /// # Errors
    ///
    /// Returns an error if the file cannot be read or the Dockerfile is malformed.
    pub fn from_file(path: impl AsRef<Path>) -> Result<Self> {
        let content =
            std::fs::read_to_string(path.as_ref()).map_err(|e| BuildError::ContextRead {
                path: path.as_ref().to_path_buf(),
                source: e,
            })?;

        Self::parse(&content)
    }

    /// Convert from the raw dockerfile-parser types to our internal representation
    fn from_raw(raw: RawDockerfile) -> Result<Self> {
        let mut global_args = Vec::new();
        let mut stages = Vec::new();
        let mut current_stage: Option<Stage> = None;
        let mut stage_index = 0;
        // Track stage names declared so far so subsequent FROM lines can
        // resolve `FROM <name>` to a stage reference even when the name
        // is also a syntactically-valid OCI reference (e.g. `FROM builder`).
        let mut known_stage_names: std::collections::HashSet<String> =
            std::collections::HashSet::new();

        for instruction in raw.instructions {
            match &instruction {
                RawInstruction::From(from) => {
                    // Save previous stage if any
                    if let Some(stage) = current_stage.take() {
                        stages.push(stage);
                    }

                    // Parse base image
                    let raw_from = from.image.content.trim().to_string();
                    let mut base_image = DockerfileFromTarget::parse(&raw_from);

                    // Post-hoc stage promotion: `DockerfileFromTarget::parse`
                    // delegates to the OCI grammar, which accepts bare names
                    // like `builder` as valid image refs. If the raw FROM
                    // text matches a previously-declared stage name, swap
                    // the parsed `Image` for a `Stage` reference.
                    if matches!(base_image, DockerfileFromTarget::Image(_))
                        && known_stage_names.contains(&raw_from)
                    {
                        base_image = DockerfileFromTarget::Stage(raw_from.clone());
                    }

                    // Get alias (stage name) - the field is `alias` not `image_alias`
                    let name = from.alias.as_ref().map(|a| a.content.clone());

                    if let Some(ref n) = name {
                        known_stage_names.insert(n.clone());
                    }

                    // Get platform flag
                    let platform = from
                        .flags
                        .iter()
                        .find(|f| f.name.content.as_str() == "platform")
                        .map(|f| f.value.to_string());

                    current_stage = Some(Stage {
                        index: stage_index,
                        name,
                        base_image,
                        platform,
                        instructions: Vec::new(),
                    });

                    stage_index += 1;
                }

                RawInstruction::Arg(arg) => {
                    let arg_inst = ArgInstruction {
                        name: arg.name.to_string(),
                        default: arg.value.as_ref().map(std::string::ToString::to_string),
                    };

                    if current_stage.is_none() {
                        global_args.push(arg_inst);
                    } else if let Some(ref mut stage) = current_stage {
                        stage.instructions.push(Instruction::Arg(arg_inst));
                    }
                }

                _ => {
                    if let Some(ref mut stage) = current_stage {
                        if let Some(inst) = Self::convert_instruction(&instruction)? {
                            stage.instructions.push(inst);
                        }
                    }
                }
            }
        }

        // Don't forget the last stage
        if let Some(stage) = current_stage {
            stages.push(stage);
        }

        // Resolve stage references in COPY --from
        // (This is currently a no-op as stage references are already correct,
        // but kept for future validation/resolution logic)
        let _stage_names: HashMap<String, usize> = stages
            .iter()
            .filter_map(|s| s.name.as_ref().map(|n| (n.clone(), s.index)))
            .collect();
        let _num_stages = stages.len();

        Ok(Self {
            global_args,
            stages,
        })
    }

    /// Convert a raw instruction to our internal representation
    #[allow(clippy::too_many_lines)]
    fn convert_instruction(raw: &RawInstruction) -> Result<Option<Instruction>> {
        let instruction = match raw {
            RawInstruction::From(_) => {
                return Ok(None);
            }

            RawInstruction::Run(run) => {
                let command = match &run.expr {
                    dockerfile_parser::ShellOrExecExpr::Shell(s) => {
                        ShellOrExec::Shell(s.to_string())
                    }
                    dockerfile_parser::ShellOrExecExpr::Exec(args) => {
                        ShellOrExec::Exec(args.elements.iter().map(|s| s.content.clone()).collect())
                    }
                };

                Instruction::Run(RunInstruction {
                    command,
                    mounts: Vec::new(),
                    network: None,
                    security: None,
                    env: HashMap::new(),
                })
            }

            RawInstruction::Copy(copy) => {
                let from = copy
                    .flags
                    .iter()
                    .find(|f| f.name.content.as_str() == "from")
                    .map(|f| f.value.to_string());

                let chown = copy
                    .flags
                    .iter()
                    .find(|f| f.name.content.as_str() == "chown")
                    .map(|f| f.value.to_string());

                let chmod = copy
                    .flags
                    .iter()
                    .find(|f| f.name.content.as_str() == "chmod")
                    .map(|f| f.value.to_string());

                let link = copy.flags.iter().any(|f| f.name.content.as_str() == "link");

                // The external parser separates sources and destination already.
                let sources: Vec<String> = copy
                    .sources
                    .iter()
                    .map(std::string::ToString::to_string)
                    .collect();
                let destination = copy.destination.to_string();

                Instruction::Copy(CopyInstruction {
                    sources,
                    destination,
                    from,
                    chown,
                    chmod,
                    link,
                    exclude: Vec::new(),
                })
            }

            RawInstruction::Entrypoint(ep) => {
                let command = match &ep.expr {
                    dockerfile_parser::ShellOrExecExpr::Shell(s) => {
                        ShellOrExec::Shell(s.to_string())
                    }
                    dockerfile_parser::ShellOrExecExpr::Exec(args) => {
                        ShellOrExec::Exec(args.elements.iter().map(|s| s.content.clone()).collect())
                    }
                };
                Instruction::Entrypoint(command)
            }

            RawInstruction::Cmd(cmd) => {
                let command = match &cmd.expr {
                    dockerfile_parser::ShellOrExecExpr::Shell(s) => {
                        ShellOrExec::Shell(s.to_string())
                    }
                    dockerfile_parser::ShellOrExecExpr::Exec(args) => {
                        ShellOrExec::Exec(args.elements.iter().map(|s| s.content.clone()).collect())
                    }
                };
                Instruction::Cmd(command)
            }

            RawInstruction::Env(env) => {
                let mut vars = HashMap::new();
                for var in &env.vars {
                    vars.insert(var.key.to_string(), var.value.to_string());
                }
                Instruction::Env(EnvInstruction { vars })
            }

            RawInstruction::Label(label) => {
                let mut labels = HashMap::new();
                for l in &label.labels {
                    labels.insert(l.name.to_string(), l.value.to_string());
                }
                Instruction::Label(labels)
            }

            RawInstruction::Arg(arg) => Instruction::Arg(ArgInstruction {
                name: arg.name.to_string(),
                default: arg.value.as_ref().map(std::string::ToString::to_string),
            }),

            RawInstruction::Misc(misc) => {
                let instruction_upper = misc.instruction.content.to_uppercase();
                match instruction_upper.as_str() {
                    "WORKDIR" => Instruction::Workdir(misc.arguments.to_string()),

                    "USER" => Instruction::User(misc.arguments.to_string()),

                    "VOLUME" => {
                        let args = misc.arguments.to_string();
                        let volumes = if args.trim().starts_with('[') {
                            serde_json::from_str(&args).unwrap_or_else(|_| vec![args])
                        } else {
                            args.split_whitespace().map(String::from).collect()
                        };
                        Instruction::Volume(volumes)
                    }

                    "EXPOSE" => {
                        let args = misc.arguments.to_string();
                        let (port_str, protocol) = if let Some((p, proto)) = args.split_once('/') {
                            let proto = match proto.to_lowercase().as_str() {
                                "udp" => ExposeProtocol::Udp,
                                _ => ExposeProtocol::Tcp,
                            };
                            (p, proto)
                        } else {
                            (args.as_str(), ExposeProtocol::Tcp)
                        };

                        let port: u16 = port_str.trim().parse().map_err(|_| {
                            BuildError::InvalidInstruction {
                                instruction: "EXPOSE".to_string(),
                                reason: format!("Invalid port number: {port_str}"),
                            }
                        })?;

                        Instruction::Expose(ExposeInstruction { port, protocol })
                    }

                    "SHELL" => {
                        let args = misc.arguments.to_string();
                        let shell: Vec<String> = serde_json::from_str(&args).map_err(|_| {
                            BuildError::InvalidInstruction {
                                instruction: "SHELL".to_string(),
                                reason: "SHELL requires a JSON array".to_string(),
                            }
                        })?;
                        Instruction::Shell(shell)
                    }

                    "STOPSIGNAL" => Instruction::Stopsignal(misc.arguments.to_string()),

                    "HEALTHCHECK" => {
                        let args = misc.arguments.to_string().trim().to_string();
                        if args.eq_ignore_ascii_case("NONE") {
                            Instruction::Healthcheck(HealthcheckInstruction::None)
                        } else {
                            let command = if let Some(stripped) = args.strip_prefix("CMD ") {
                                ShellOrExec::Shell(stripped.to_string())
                            } else {
                                ShellOrExec::Shell(args)
                            };
                            Instruction::Healthcheck(HealthcheckInstruction::cmd(command))
                        }
                    }

                    "ONBUILD" => {
                        tracing::warn!("ONBUILD instruction parsing not fully implemented");
                        return Ok(None);
                    }

                    "MAINTAINER" => {
                        let mut labels = HashMap::new();
                        labels.insert("maintainer".to_string(), misc.arguments.to_string());
                        Instruction::Label(labels)
                    }

                    "ADD" => {
                        let args = misc.arguments.to_string();
                        let parts: Vec<String> =
                            args.split_whitespace().map(String::from).collect();

                        if parts.len() < 2 {
                            return Err(BuildError::InvalidInstruction {
                                instruction: "ADD".to_string(),
                                reason: "ADD requires at least one source and a destination"
                                    .to_string(),
                            });
                        }

                        let (sources, dest) = parts.split_at(parts.len() - 1);
                        let destination = dest.first().cloned().unwrap_or_default();

                        Instruction::Add(AddInstruction {
                            sources: sources.to_vec(),
                            destination,
                            chown: None,
                            chmod: None,
                            link: false,
                            checksum: None,
                            keep_git_dir: false,
                        })
                    }

                    other => {
                        tracing::warn!("Unknown Dockerfile instruction: {}", other);
                        return Ok(None);
                    }
                }
            }
        };

        Ok(Some(instruction))
    }

    /// Get a stage by name or index
    #[must_use]
    pub fn get_stage(&self, name_or_index: &str) -> Option<&Stage> {
        self.stages.iter().find(|s| s.matches(name_or_index))
    }

    /// Get the final stage (last one in the Dockerfile)
    #[must_use]
    pub fn final_stage(&self) -> Option<&Stage> {
        self.stages.last()
    }

    /// Get all stage names/identifiers
    #[must_use]
    pub fn stage_names(&self) -> Vec<String> {
        self.stages.iter().map(Stage::identifier).collect()
    }

    /// Check if a stage exists
    #[must_use]
    pub fn has_stage(&self, name_or_index: &str) -> bool {
        self.get_stage(name_or_index).is_some()
    }

    /// Returns the number of stages
    #[must_use]
    pub fn stage_count(&self) -> usize {
        self.stages.len()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_parse_simple_dockerfile() {
        let content = r#"
FROM alpine:3.18
RUN apk add --no-cache curl
COPY . /app
WORKDIR /app
CMD ["./app"]
"#;

        let dockerfile = Dockerfile::parse(content).unwrap();
        assert_eq!(dockerfile.stages.len(), 1);

        let stage = &dockerfile.stages[0];
        assert_eq!(stage.index, 0);
        assert!(stage.name.is_none());
        assert_eq!(stage.instructions.len(), 4);
    }

    #[test]
    fn test_parse_multistage_dockerfile() {
        let content = r#"
FROM golang:1.21 AS builder
WORKDIR /src
COPY . .
RUN go build -o /app

FROM alpine:3.18
COPY --from=builder /app /app
CMD ["/app"]
"#;

        let dockerfile = Dockerfile::parse(content).unwrap();
        assert_eq!(dockerfile.stages.len(), 2);

        let builder = &dockerfile.stages[0];
        assert_eq!(builder.name, Some("builder".to_string()));

        let runtime = &dockerfile.stages[1];
        assert!(runtime.name.is_none());

        let copy = runtime
            .instructions
            .iter()
            .find(|i| matches!(i, Instruction::Copy(_)));
        assert!(copy.is_some());
        if let Some(Instruction::Copy(c)) = copy {
            assert_eq!(c.from, Some("builder".to_string()));
        }
    }

    #[test]
    fn test_parse_copy_from_external_image_reference() {
        // `COPY --from=<external-image>` must capture the full registry-
        // qualified reference in `CopyInstruction.from` so the buildah
        // backend can pull and forward it to `buildah copy --from=...`.
        let content = r"
FROM alpine:3.18
COPY --from=ghcr.io/astral-sh/uv:0.5.0 /uv /usr/local/bin/uv
RUN /usr/local/bin/uv --version
";

        let dockerfile = Dockerfile::parse(content).unwrap();
        assert_eq!(dockerfile.stages.len(), 1);

        let copy = dockerfile.stages[0]
            .instructions
            .iter()
            .find_map(|i| {
                if let Instruction::Copy(c) = i {
                    Some(c)
                } else {
                    None
                }
            })
            .expect("COPY instruction present");

        assert_eq!(
            copy.from,
            Some("ghcr.io/astral-sh/uv:0.5.0".to_string()),
            "external image ref must be preserved verbatim in CopyInstruction.from",
        );
        assert_eq!(copy.sources, vec!["/uv".to_string()]);
        assert_eq!(copy.destination, "/usr/local/bin/uv".to_string());

        // The parser must NOT treat the external ref as a stage; only the
        // top-level `FROM alpine:3.18` should appear in the stage list.
        assert!(dockerfile.get_stage("ghcr.io/astral-sh/uv:0.5.0").is_none());
    }

    #[test]
    fn expand_from_args_all_forms() {
        let vars: HashMap<String, String> = [
            ("BASE".to_string(), "ghcr.io/org/img".to_string()),
            ("TAG".to_string(), "1.2".to_string()),
            ("EMPTY".to_string(), String::new()),
        ]
        .into_iter()
        .collect();

        assert_eq!(
            expand_from_args("${BASE}:${TAG}", &vars),
            "ghcr.io/org/img:1.2"
        );
        assert_eq!(expand_from_args("$BASE", &vars), "ghcr.io/org/img");
        assert_eq!(
            expand_from_args("img:${MISSING:-latest}", &vars),
            "img:latest"
        );
        assert_eq!(
            expand_from_args("img:${EMPTY:-fallback}", &vars),
            "img:fallback"
        );
        assert_eq!(expand_from_args("img:${TAG:+pinned}", &vars), "img:pinned");
        assert_eq!(expand_from_args("img:${MISSING:+pinned}", &vars), "img:");
        assert_eq!(expand_from_args("${MISSING}", &vars), "");
        // BuildKit `$$` escape yields a literal dollar; a trailing `$`
        // passes through.
        assert_eq!(expand_from_args("a$$b", &vars), "a$b");
        assert_eq!(expand_from_args("price$", &vars), "price$");
    }

    #[test]
    fn resolve_from_args_expands_from_lines() {
        let content = r"
ARG BASE_IMAGE=ghcr.io/org/alpine:latest
ARG BASE_TAG=latest
FROM ${BASE_IMAGE} AS builder
RUN echo hi
FROM ghcr.io/org/alpine:${BASE_TAG}
COPY --from=builder /x /x
";
        let mut dockerfile = Dockerfile::parse(content).unwrap();
        // Pre-resolution both FROM targets are (mis)classified as stages.
        assert!(dockerfile.stages[0].base_image.is_stage());
        assert!(dockerfile.stages[1].base_image.is_stage());

        let build_args: HashMap<String, String> = [("BASE_TAG".to_string(), "3.20".to_string())]
            .into_iter()
            .collect();
        dockerfile.resolve_from_args(&build_args);

        match &dockerfile.stages[0].base_image {
            DockerfileFromTarget::Image(r) => {
                assert_eq!(r.to_string(), "ghcr.io/org/alpine:latest");
            }
            other => panic!("stage 0 not resolved to an image: {other:?}"),
        }
        match &dockerfile.stages[1].base_image {
            DockerfileFromTarget::Image(r) => {
                // The build arg overrides the declared default.
                assert_eq!(r.to_string(), "ghcr.io/org/alpine:3.20");
            }
            other => panic!("stage 1 not resolved to an image: {other:?}"),
        }
    }

    #[test]
    fn resolve_from_args_keeps_stage_references() {
        let content = r"
ARG BASE=ghcr.io/org/alpine:latest
FROM ${BASE} AS builder
RUN echo hi
FROM builder
RUN echo again
";
        let mut dockerfile = Dockerfile::parse(content).unwrap();
        dockerfile.resolve_from_args(&HashMap::new());
        assert!(matches!(
            &dockerfile.stages[0].base_image,
            DockerfileFromTarget::Image(_)
        ));
        // `FROM builder` must stay a stage reference.
        assert_eq!(
            dockerfile.stages[1].base_image,
            DockerfileFromTarget::Stage("builder".to_string())
        );
    }

    #[test]
    fn test_parse_global_args() {
        let content = r#"
ARG BASE_IMAGE=alpine:3.18
FROM ${BASE_IMAGE}
RUN echo "hello"
"#;

        let dockerfile = Dockerfile::parse(content).unwrap();
        assert_eq!(dockerfile.global_args.len(), 1);
        assert_eq!(dockerfile.global_args[0].name, "BASE_IMAGE");
        assert_eq!(
            dockerfile.global_args[0].default,
            Some("alpine:3.18".to_string())
        );
    }

    #[test]
    fn test_get_stage_by_name() {
        let content = r#"
FROM alpine:3.18 AS base
RUN echo "base"

FROM base AS builder
RUN echo "builder"
"#;

        let dockerfile = Dockerfile::parse(content).unwrap();

        let base = dockerfile.get_stage("base");
        assert!(base.is_some());
        assert_eq!(base.unwrap().index, 0);

        let builder = dockerfile.get_stage("builder");
        assert!(builder.is_some());
        assert_eq!(builder.unwrap().index, 1);

        let stage_0 = dockerfile.get_stage("0");
        assert!(stage_0.is_some());
        assert_eq!(stage_0.unwrap().name, Some("base".to_string()));
    }

    #[test]
    fn test_final_stage() {
        let content = r#"
FROM alpine:3.18 AS builder
RUN echo "builder"

FROM scratch
COPY --from=builder /app /app
"#;

        let dockerfile = Dockerfile::parse(content).unwrap();
        let final_stage = dockerfile.final_stage().unwrap();

        assert_eq!(final_stage.index, 1);
        assert!(matches!(
            final_stage.base_image,
            DockerfileFromTarget::Scratch
        ));
    }

    #[test]
    fn test_parse_env_instruction() {
        let content = r"
FROM alpine
ENV FOO=bar BAZ=qux
";

        let dockerfile = Dockerfile::parse(content).unwrap();
        let stage = &dockerfile.stages[0];

        let env = stage
            .instructions
            .iter()
            .find(|i| matches!(i, Instruction::Env(_)));
        assert!(env.is_some());

        if let Some(Instruction::Env(e)) = env {
            assert_eq!(e.vars.get("FOO"), Some(&"bar".to_string()));
            assert_eq!(e.vars.get("BAZ"), Some(&"qux".to_string()));
        }
    }
}