pdk-test 1.7.0

// Copyright (c) 2026, Salesforce, Inc.,
// All rights reserved.
// For full license text, see the LICENSE.txt file

use std::collections::HashMap;

use bollard::models::{PortBinding, PortMap};
use bollard::service::{Mount, MountTypeEnum};

use crate::error::TestError;
use crate::host::Mode;
use crate::image::Image;
use crate::port::{Port, PortAccess};
use crate::probe::Probe;

/// Type alias for sockets mapping between container ports and hosts.
pub(crate) type Sockets = HashMap<Port, String>;

/// Struct to manage the container config.
#[derive(Debug)]
pub struct ContainerConfig {
    image: Image,
    hostname: String,
    ports: Vec<PortAccess>,
    mounts: Vec<(String, String, String)>,
    env: Vec<(String, String)>,
    readiness: Option<Box<dyn Probe>>,
    cmd: Vec<String>,
}

impl ContainerConfig {
    /// Create a new container config from a hostname and image.
    pub fn new(hostname: String, image: Image) -> Self {
        Self {
            image,
            hostname,
            ports: vec![],
            mounts: vec![],
            env: vec![],
            readiness: None,
            cmd: vec![],
        }
    }

    /// Get the container hostname.
    pub fn hostname(&self) -> &str {
        &self.hostname
    }

    /// Get the container image.
    pub fn image(&self) -> &Image {
        &self.image
    }

    /// Get the container ports.
    pub fn ports(&self) -> &[PortAccess] {
        &self.ports
    }

    /// Get the container mounts.
    pub fn mounts(&self) -> &[(String, String, String)] {
        &self.mounts
    }

    /// Get the container environment variables.
    pub fn env(&self) -> &[(String, String)] {
        &self.env
    }

    /// Get the container readiness probe.
    pub fn readiness(&self) -> Option<&dyn Probe> {
        self.readiness.as_deref()
    }

    /// Get the container command.
    pub fn cmd(&self) -> &[String] {
        &self.cmd
    }

    /// Create a new container config builder.
    pub fn builder<H>(hostname: H, image: Image) -> ContainerConfigBuilder
    where
        H: Into<String>,
    {
        ContainerConfigBuilder::new(hostname, image)
    }

    /// Get the container formatted environment variables.
    pub(super) fn formatted_env(&self) -> Vec<String> {
        self.env
            .iter()
            .map(|(var, value)| format!("{var}={value}"))
            .collect()
    }

    fn standalone_pas(&self) -> Result<(PortMap, Sockets), TestError> {
        let bindings: Vec<_> = self
            .ports
            .iter()
            .filter(|p| p.visibility().is_published())
            .map(|p| {
                p.public_port()
                    .map(|e| (p.port(), e))
                    .ok_or(TestError::UnavailablePorts)
            })
            .collect::<Result<_, _>>()?;

        let sockets = bindings
            .iter()
            .map(|(internal, external)| (*internal, format!("localhost:{external}")))
            .collect();

        let port_bindings = bindings
            .iter()
            .map(|(internal, external)| {
                (
                    format!("{internal}/tcp"),
                    Some(vec![PortBinding {
                        host_ip: None,
                        host_port: Some(external.to_string()),
                    }]),
                )
            })
            .collect();

        Ok((port_bindings, sockets))
    }

    fn container_pas(&self) -> Result<(PortMap, Sockets), TestError> {
        let hostname = &self.hostname;
        let sockets = self
            .ports
            .iter()
            .filter(|p| p.visibility().is_published())
            .map(PortAccess::port)
            .map(|port| (port, format!("{hostname}:{port}")))
            .collect();
        Ok((PortMap::new(), sockets))
    }

    /// Get the container ports and sockets.
    pub(super) fn ports_and_sockets(&self, mode: Mode) -> Result<(PortMap, Sockets), TestError> {
        match mode {
            Mode::Standalone => self.standalone_pas(),
            Mode::Containerized => self.container_pas(),
        }
    }

    /// Get the container Docker mounts.
    pub fn docker_mounts(&self) -> Vec<Mount> {
        self.mounts()
            .iter()
            .map(|(src, dst_base, dst)| bollard::models::Mount {
                typ: Some(MountTypeEnum::BIND),
                source: Some(src.to_string()),
                target: Some(format!("/{dst_base}/{dst}")),
                ..Default::default()
            })
            .collect()
    }
}

/// Builder for initializing a [`ContainerConfig`].
#[derive(Debug)]
pub struct ContainerConfigBuilder {
    config: ContainerConfig,
}

impl ContainerConfigBuilder {
    fn new<H>(hostname: H, image: Image) -> Self
    where
        H: Into<String>,
    {
        Self {
            config: ContainerConfig::new(hostname.into(), image),
        }
    }

    /// Set the container ports.
    pub fn ports<T>(self, ports: T) -> Self
    where
        T: IntoIterator<Item = PortAccess>,
    {
        Self {
            config: ContainerConfig {
                ports: ports.into_iter().collect(),
                ..self.config
            },
        }
    }

    /// Set the container mounts.
    pub fn mounts<T, S, B, D>(self, mounts: T) -> Self
    where
        T: IntoIterator<Item = (S, B, D)>,
        S: Into<String>,
        B: Into<String>,
        D: Into<String>,
    {
        Self {
            config: ContainerConfig {
                mounts: mounts
                    .into_iter()
                    .map(|(s, b, d)| (s.into(), b.into(), d.into()))
                    .collect(),
                ..self.config
            },
        }
    }

    /// Set the container environment variables.
    pub fn env<T, N, V>(self, env: T) -> Self
    where
        T: IntoIterator<Item = (N, V)>,
        N: Into<String>,
        V: ToString,
    {
        Self {
            config: ContainerConfig {
                env: env
                    .into_iter()
                    .map(|(n, v)| (n.into(), v.to_string()))
                    .collect(),
                ..self.config
            },
        }
    }

    /// Set the container readiness probe.
    pub fn readiness<T: Probe + 'static>(self, readiness: T) -> Self {
        Self {
            config: ContainerConfig {
                readiness: Some(Box::new(readiness)),
                ..self.config
            },
        }
    }

    /// Set the container command.
    pub fn cmd<T, C>(self, cmd: T) -> Self
    where
        T: IntoIterator<Item = C>,
        C: Into<String>,
    {
        Self {
            config: ContainerConfig {
                cmd: cmd.into_iter().map(|c| c.into()).collect(),
                ..self.config
            },
        }
    }

    /// Build the container config.
    pub fn build(self) -> ContainerConfig {
        self.config
    }
}

/// Trait for configuring a container.
pub trait Config {
    fn hostname(&self) -> &str;

    fn port(&self) -> Port;

    fn schema(&self) -> &str;

    fn to_container_config(&self) -> Result<ContainerConfig, TestError>;
}