#!/usr/bin/env python3

# author: Ole Schuett

import argparse
import io
from pathlib import Path
from typing import Any


# ======================================================================================
def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--check", action="store_true")
    args = parser.parse_args()

    for version in "sdbg", "ssmp", "pdbg", "psmp":
        with OutputFile(f"Dockerfile.test_{version}", args.check) as f:
            mpi_mode = "mpich" if version.startswith("p") else "no"
            f.write(install_deps_toolchain(mpi_mode=mpi_mode))
            f.write(regtest("toolchain", version))

    with OutputFile(f"Dockerfile.test_psmp_4ranks", args.check) as f:
        testopts = f"--ompthreads=1 --mpiranks=4"
        f.write(install_deps_toolchain())
        f.write(regtest("toolchain", "psmp", testopts=testopts))

    with OutputFile(f"Dockerfile.test_generic_psmp", args.check) as f:
        f.write(
            install_deps_toolchain(
                target_cpu="generic", with_gauxc="no", with_libtorch="no"
            )
        )
        f.write(regtest("toolchain_generic", "psmp"))

    with OutputFile(f"Dockerfile.test_openmpi-psmp", args.check) as f:
        f.write(install_deps_toolchain(mpi_mode="openmpi"))
        f.write(regtest("toolchain", "psmp"))

    with OutputFile(f"Dockerfile.test_fedora-psmp", args.check) as f:
        f.write(install_deps_toolchain(base_image="fedora:44"))
        f.write(regtest("toolchain", "psmp"))

    with OutputFile(f"Dockerfile.test_rawhide-pdbg", args.check) as f:
        f.write(install_deps_toolchain(base_image="fedora:rawhide"))
        f.write(regtest("toolchain", "pdbg"))

    for version in "ssmp", "psmp":
        mpi_mode = "intelmpi" if version.startswith("p") else "no"
        f.write(install_deps_toolchain(mpi_mode=mpi_mode))
        with OutputFile(f"Dockerfile.test_intel-ifort-{version}", args.check) as f:
            base_image = "intel/hpckit:2024.2.1-0-devel-ubuntu22.04"
            f.write(install_deps_toolchain_intel(base_image, mpi_mode, with_ifx="no"))
            f.write(regtest("toolchain_intel", version))
        with OutputFile(f"Dockerfile.test_intel-ifx-{version}", args.check) as f:
            base_image = "intel/oneapi-hpckit:2025.2.2-0-devel-ubuntu24.04"
            f.write(install_deps_toolchain_intel(base_image, mpi_mode, with_ifx="yes"))
            f.write(regtest("toolchain_intel", version))

    with OutputFile(f"Dockerfile.test_minimal", args.check) as f:
        f.write(install_deps_ubuntu())
        f.write(regtest("minimal", "sdbg"))

    # Spack/CMake based testers

    testopts = f"--keepalive"

    with OutputFile(f"Dockerfile.test_spack_pdbg", args.check) as f:
        f.write(
            install_cp2k_spack(
                version="pdbg",
                mpi_mode="mpich",
                feature_flags="",
                testopts="",
                image_tag=f.image_tag,
            )
        )

    for gcc_version in 10, 11, 12, 13, 14, 15, 16:
        with OutputFile(
            f"Dockerfile.test_spack_psmp-gcc{gcc_version}", args.check
        ) as f:
            base_image = "ubuntu:26.04" if gcc_version > 12 else "ubuntu:24.04"
            feature_flags = "" if gcc_version > 10 else "-df libtorch"
            f.write(
                install_cp2k_spack(
                    version="psmp",
                    mpi_mode="mpich",
                    gcc_version=gcc_version,
                    base_image=base_image,
                    feature_flags=feature_flags,
                    testopts=testopts,
                    image_tag=f.image_tag,
                )
            )

    with OutputFile(f"Dockerfile.test_spack_pdbg-rawhide", args.check) as f:
        f.write(
            install_cp2k_spack(
                version="pdbg",
                mpi_mode="mpich",
                base_image="fedora:rawhide",
                feature_flags="",
                testopts=testopts,
                image_tag=f.image_tag,
            )
        )

    with OutputFile(f"Dockerfile.test_spack_psmp-fedora", args.check) as f:
        f.write(
            install_cp2k_spack(
                version="psmp",
                mpi_mode="mpich",
                base_image="fedora:latest",
                feature_flags="",
                testopts=testopts,
                image_tag=f.image_tag,
            )
        )

    with OutputFile(f"Dockerfile.test_spack_psmp-opensuse", args.check) as f:
        f.write(
            install_cp2k_spack(
                version="psmp",
                mpi_mode="mpich",
                base_image="opensuse/leap:16.0",
                gcc_version=13,
                feature_flags="",
                testopts=testopts,
                image_tag=f.image_tag,
            )
        )

    with OutputFile(f"Dockerfile.test_spack_psmp-rockylinux", args.check) as f:
        f.write(
            install_cp2k_spack(
                version="psmp",
                mpi_mode="mpich",
                base_image="docker.io/rockylinux/rockylinux:10",
                gcc_version=14,
                feature_flags="",
                testopts=testopts,
                image_tag=f.image_tag,
            )
        )

    with OutputFile(f"Dockerfile.test_spack_psmp-4x2", args.check) as f:
        f.write(
            install_cp2k_spack(
                version="psmp",
                mpi_mode="mpich",
                feature_flags="",
                testopts=f"--keepalive --mpiranks=4 --ompthreads=2",
                image_tag=f.image_tag,
            )
        )

    with OutputFile(f"Dockerfile.test_spack_openmpi-pdbg", args.check) as f:
        f.write(
            install_cp2k_spack(
                version="pdbg",
                mpi_mode="openmpi",
                feature_flags="",
                testopts="",
                image_tag=f.image_tag,
            )
        )

    with OutputFile(f"Dockerfile.test_spack_openmpi-psmp", args.check) as f:
        f.write(
            install_cp2k_spack(
                version="psmp",
                mpi_mode="openmpi",
                feature_flags="",
                testopts=testopts,
                image_tag=f.image_tag,
            )
        )

    with OutputFile(f"Dockerfile.test_spack_sdbg", args.check) as f:
        f.write(install_cp2k_spack(version="sdbg", mpi_mode="no", testopts=""))

    with OutputFile(f"Dockerfile.test_spack_ssmp", args.check) as f:
        f.write(
            install_cp2k_spack(
                version="ssmp", mpi_mode="no", testopts=testopts, image_tag=f.image_tag
            )
        )

    with OutputFile(f"Dockerfile.test_spack_ssmp-static", args.check) as f:
        f.write(
            install_cp2k_spack(
                version="ssmp-static",
                mpi_mode="no",
                testopts=testopts,
                image_tag=f.image_tag,
            )
        )

    with OutputFile(f"Dockerfile.test_spack_ssmp-P100", args.check) as f:
        f.write(
            install_cp2k_spack(
                version="ssmp",
                mpi_mode="no",
                base_image="docker.io/nvidia/cuda:12.9.1-devel-ubuntu24.04",
                gcc_version=13,
                gpu_model="P100",
                testopts=testopts,
                image_tag=f.image_tag,
            )
        )

    with OutputFile(f"Dockerfile.test_spack_psmp-P100", args.check) as f:
        f.write(
            install_cp2k_spack(
                version="psmp",
                mpi_mode="mpich",
                base_image="docker.io/nvidia/cuda:12.9.1-devel-ubuntu24.04",
                gcc_version=13,
                gpu_model="P100",
                feature_flags="",
                testopts=testopts,
                image_tag=f.image_tag,
            )
        )

    with OutputFile(f"Dockerfile.test_spack_pdbg-P100", args.check) as f:
        f.write(
            install_cp2k_spack(
                version="pdbg",
                mpi_mode="mpich",
                base_image="docker.io/nvidia/cuda:12.9.1-devel-ubuntu24.04",
                gcc_version=13,
                gpu_model="P100",
                feature_flags="",
                testopts=testopts,
                image_tag=f.image_tag,
            )
        )

    # End Spack/CMake based tester

    with OutputFile(f"Dockerfile.test_asan-psmp", args.check) as f:
        f.write(install_deps_toolchain())
        f.write(regtest("toolchain_asan", "psmp"))

    with OutputFile(f"Dockerfile.test_coverage", args.check) as f:
        f.write(install_deps_toolchain())
        f.write(coverage())

    for gcc_version in 9, 10, 11, 12, 13, 14, 15, 16:
        with OutputFile(f"Dockerfile.test_gcc{gcc_version}", args.check) as f:
            f.write(install_deps_ubuntu(gcc_version=gcc_version))
            f.write(regtest("ubuntu", "ssmp"))

    with OutputFile("Dockerfile.test_arm64-psmp", args.check) as f:
        base_img = "arm64v8/ubuntu:26.04"
        libs = dict(with_libtorch="no", with_deepmd="no", with_gauxc="no")
        f.write(install_deps_toolchain(base_img, **libs))
        f.write(regtest("toolchain_arm64", "psmp"))

    with OutputFile(f"Dockerfile.test_performance", args.check) as f:
        f.write(install_deps_toolchain())
        f.write(performance("toolchain"))

    for gpu_ver in "P100", "V100", "A100":
        with OutputFile(f"Dockerfile.test_cuda_{gpu_ver}", args.check) as f:
            f.write(install_deps_toolchain_cuda(gpu_ver=gpu_ver))
            f.write(regtest(f"toolchain_cuda_{gpu_ver}", "psmp"))
        with OutputFile(f"Dockerfile.test_performance_cuda_{gpu_ver}", args.check) as f:
            f.write(install_deps_toolchain_cuda(gpu_ver=gpu_ver))
            f.write(performance(f"toolchain_cuda_{gpu_ver}"))

    for gpu_ver in "Mi50", "Mi100":
        with OutputFile(f"Dockerfile.build_hip_rocm_{gpu_ver}", args.check) as f:
            f.write(install_deps_toolchain_hip_rocm(gpu_ver=gpu_ver))
            f.write(test_build(f"toolchain_hip_{gpu_ver}", "psmp"))

    with OutputFile(f"Dockerfile.test_conventions", args.check) as f:
        f.write(install_deps_toolchain())
        f.write(conventions())

    with OutputFile(f"Dockerfile.test_manual", args.check) as f:
        f.write(install_deps_toolchain())
        f.write(manual())

    with OutputFile(f"Dockerfile.test_precommit", args.check) as f:
        f.write(precommit())

    for name in "ase", "aiida", "i-pi", "phonopy", "gromacs":
        with OutputFile(f"Dockerfile.test_{name}", args.check) as f:
            f.write(install_deps_toolchain(mpi_mode="no"))
            f.write(test_3rd_party(name))

    for name in "misc", "doxygen":
        with OutputFile(f"Dockerfile.test_{name}", args.check) as f:
            f.write(test_without_build(name))


# ======================================================================================
def regtest(profile: str, version: str, testopts: str = "") -> str:
    return install_cp2k(profile=profile, version=version) + rf"""
# Run regression tests.
ARG TESTOPTS="{testopts}"
COPY ./tests ./tests
COPY ./tools/docker/scripts/test_regtest.sh ./
RUN /bin/bash -o pipefail -c " \
    TESTOPTS='${{TESTOPTS}}' \
    ./test_regtest.sh {profile} {version} |& tee report.log && \
    rm -rf regtesting"
""" + print_cached_report()


# ======================================================================================
def test_build(profile: str, version: str) -> str:
    return install_cp2k(profile=profile, version=version) + rf"""
# Run build test.
COPY ./tools/docker/scripts/test_build.sh .
RUN ./test_build.sh "{profile}" "{version}" 2>&1 | tee report.log
""" + print_cached_report()


# ======================================================================================
def performance(profile: str) -> str:
    return install_cp2k(profile=profile, version="psmp") + rf"""
# Run performance test for {profile}.
COPY ./benchmarks ./benchmarks
COPY ./tools/regtesting ./tools/regtesting
COPY ./tools/docker/scripts/test_performance.sh  \
     ./tools/docker/scripts/plot_performance.py  \
     ./
RUN ./test_performance.sh "{profile}" 2>&1 | tee report.log
""" + print_cached_report()


# ======================================================================================
def coverage() -> str:
    return (
        install_cp2k(profile="toolchain_coverage", version="psmp", revision=True)
        + rf"""
# Run coverage test.
COPY ./tests ./tests
COPY ./tools/docker/scripts/test_coverage.sh .
RUN ./test_coverage.sh 2>&1 | tee report.log
"""
        + print_cached_report()
    )


# ======================================================================================
def conventions() -> str:
    return (
        f"""
COPY ./tools/conventions/redirect_gfortran_output.py /usr/bin/
"""
        + install_cp2k(profile="toolchain_conventions", version="psmp")
        + f"""
# Run test for conventions.
COPY ./tools/conventions ./tools/conventions
RUN /bin/bash -ec "./tools/conventions/test_conventions.sh |& tee report.log"
"""
        + print_cached_report()
    )


# ======================================================================================
def manual() -> str:
    return install_cp2k(profile="toolchain", version="pdbg", revision=True) + rf"""
# Generate manual.
COPY ./docs ./docs
COPY ./tools/input_editing ./tools/input_editing
COPY ./tools/docker/scripts/test_manual.sh .
ARG ADD_EDIT_LINKS=yes
RUN ./test_manual.sh "${{ADD_EDIT_LINKS}}" 2>&1 | tee report.log
""" + print_cached_report()


# ======================================================================================
def precommit() -> str:
    return rf"""
FROM ubuntu:24.04

# Install dependencies.
WORKDIR /opt/cp2k-precommit
COPY ./tools/precommit/ /opt/cp2k-precommit/
RUN ./install_requirements.sh
ENV PATH="/opt/venv/bin:/opt/cp2k-precommit:$PATH"

# Install sources.
WORKDIR /opt/cp2k
COPY ./ ./

# Run precommit test.
RUN ./tools/docker/scripts/test_precommit.sh 2>&1 | tee report.log
""" + print_cached_report()


# ======================================================================================
def test_3rd_party(name: str) -> str:
    return install_cp2k(profile="toolchain", version="ssmp") + rf"""
# Run test for {name}.
COPY ./tests ./tests
COPY ./tools/docker/scripts/test_{name}.sh ./
RUN ./test_{name}.sh 2>&1 | tee report.log
""" + print_cached_report()


# ======================================================================================
def test_without_build(name: str) -> str:
    return rf"""
FROM ubuntu:24.04

# Install dependencies.
WORKDIR /opt/cp2k
COPY ./tools/docker/scripts/install_{name}.sh .
COPY ./tools/pao-ml/requirements.txt pao-ml-requirements.txt
RUN ./install_{name}.sh
ENV PATH="/opt/venv/bin:$PATH"

# Install sources.
ARG GIT_COMMIT_SHA
COPY ./src ./src
COPY ./data ./data
COPY ./docs ./docs
COPY ./tools ./tools
COPY ./tests ./tests
COPY ./cmake ./cmake
COPY ./CMakeLists.txt .
RUN bash -c "if [ -n "${{GIT_COMMIT_SHA}}" ] ; then echo "git:\${{GIT_COMMIT_SHA::7}}" > REVISION; fi"

# Run test for {name}.
COPY ./tools/docker/scripts/test_{name}.sh .
RUN ./test_{name}.sh 2>&1 | tee report.log
""" + print_cached_report()


# ======================================================================================
def print_cached_report() -> str:
    return r"""
# Output the report if the image is old and was therefore pulled from the build cache.
CMD cat $(find ./report.log -mmin +10) | sed '/^Summary:/ s/$/ (cached)/'
ENTRYPOINT []

#EOF
"""


# ======================================================================================
def install_cp2k(profile: str, version: str, revision: bool = False) -> str:
    output = ""
    if revision:
        output += "\n"
        output += "ARG GIT_COMMIT_SHA\n"
        output += "ENV GIT_COMMIT_SHA=${GIT_COMMIT_SHA}\n"

    output += rf"""
# Install CP2K sources.
WORKDIR /opt/cp2k
COPY ./src ./src
COPY ./data ./data
COPY ./tools/build_utils ./tools/build_utils
COPY ./cmake ./cmake
COPY ./CMakeLists.txt .

# Compile CP2K.
COPY ./tools/docker/scripts/build_cp2k.sh ./tools/docker/scripts/cmake_cp2k.sh ./
RUN ./build_cp2k.sh {profile} {version}
"""
    return output


# ======================================================================================
def install_deps_toolchain(
    base_image: str = "ubuntu:26.04",
    mpi_mode: str = "mpich",
    with_dbcsr: str = "",  # enabled by default
    with_gcc: str = "system",
    **kwargs: str,
) -> str:
    output = f"\nFROM {base_image}\n\n"
    output += install_toolchain(
        base_image=base_image,
        install_all="",
        mpi_mode=mpi_mode,
        with_dbcsr=with_dbcsr,
        with_gcc=with_gcc,
        **kwargs,
    )
    return output


# ======================================================================================
def install_deps_ubuntu(gcc_version: int = 15) -> str:
    assert gcc_version > 8
    base_image = "ubuntu:26.04" if gcc_version > 14 else "ubuntu:24.04"
    output = f"\nFROM {base_image}\n"

    if gcc_version > 13:
        output += rf"""
# Add Ubuntu universe repository.
RUN apt-get update -qq && apt-get install -qq --no-install-recommends software-properties-common
RUN add-apt-repository universe
"""

    output += rf"""
# Install Ubuntu packages.
RUN export DEBIAN_FRONTEND=noninteractive DEBCONF_NONINTERACTIVE_SEEN=true && \
    apt-get update -qq && apt-get install -qq --no-install-recommends \
    cmake \
    less \
    nano \
    make \
    ninja-build \
    wget \
    python3 \
    ca-certificates \
    gcc-{gcc_version} \
    g++-{gcc_version} \
    gfortran-{gcc_version} \
    libfftw3-dev \
    libopenblas-dev \
    libint2-dev \
    libxc-dev \
    libhdf5-dev \
    libxsmm-dev \
    libspglib-f08-dev \
   && rm -rf /var/lib/apt/lists/*

# Create links in /usr/local/bin to overrule links in /usr/bin.
RUN ln -sf /usr/bin/gcc-{gcc_version}      /usr/local/bin/gcc  && \
    ln -sf /usr/bin/g++-{gcc_version}      /usr/local/bin/g++  && \
    ln -sf /usr/bin/gfortran-{gcc_version} /usr/local/bin/gfortran

# Use toolchain to install DBCSR.
""" + install_toolchain(
        base_image=base_image,
        mpi_mode="no",
        with_dbcsr="install",
        with_gcc="system",
        with_cmake="system",
        with_ninja="system",
        with_openblas="system",
        with_libxc="no",
        with_libint="no",
        with_fftw="no",
        with_libxsmm="install",
        with_libxs="install",
        with_spglib="no",
        with_libvori="no",
        with_tblite="no",
    )

    return output


# ======================================================================================
def install_deps_toolchain_intel(base_image: str, mpi_mode: str, with_ifx: str) -> str:
    return rf"""
FROM {base_image}

""" + install_toolchain(
        base_image="ubuntu",
        install_all="",
        mpi_mode=mpi_mode,
        with_ifx=with_ifx,
        with_mkl="",
        with_libsmeagol="",
        with_libtorch="no",
        with_deepmd="no",
        with_gauxc="no",
    )


# ======================================================================================
def install_deps_toolchain_cuda(gpu_ver: str, **kwargs: str) -> str:
    deps = rf"""
FROM nvidia/cuda:12.9.1-devel-ubuntu24.04

# Setup CUDA environment.
ENV CUDA_PATH /usr/local/cuda
ENV LD_LIBRARY_PATH /usr/local/cuda/lib64

# Disable JIT cache as there seems to be an issue with file locking on overlayfs.
# See also https://github.com/cp2k/cp2k/pull/2337
ENV CUDA_CACHE_DISABLE 1

# Install Ubuntu packages.
RUN apt-get update -qq && apt-get install -qq --no-install-recommends \
    gfortran                                                          \
   && rm -rf /var/lib/apt/lists/*

""" + install_toolchain(
        base_image="ubuntu",
        with_mpich="install",
        mpi_mode="mpich",
        enable_cuda="yes",
        with_libgint="install",
        with_sirius="install",
        gpu_ver=gpu_ver,
        **kwargs,
    )
    return deps


# ======================================================================================
def install_deps_toolchain_hip_rocm(gpu_ver: str) -> str:
    return rf"""
FROM rocm/dev-ubuntu-24.04:7.2-complete

# Install some Ubuntu packages.
RUN apt-get update -qq && apt-get install -qq --no-install-recommends \
    hipblas                                                           \
    gfortran                                                          \
    mpich                                                             \
    libmpich-dev                                                      \
   && rm -rf /var/lib/apt/lists/*

# Remove LTO from Ubuntu's MPICH
RUN sed -i -e 's/-flto=auto//g' -e 's/-ffat-lto-objects//g' \
    /usr/lib/x86_64-linux-gnu/pkgconfig/mpich.pc \
    /usr/bin/*.mpich

# Setup HIP environment.
ENV ROCM_PATH /opt/rocm
ENV PATH ${{PATH}}:${{ROCM_PATH}}/bin
ENV LD_LIBRARY_PATH ${{LD_LIBRARY_PATH}}:${{ROCM_PATH}}/lib
ENV HIP_PLATFORM amd
RUN hipconfig

""" + install_toolchain(
        base_image="ubuntu",
        mpi_mode="mpich",
        enable_hip="yes",
        with_sirius="install",
        gpu_ver=gpu_ver,
    )


# ======================================================================================
def install_toolchain(base_image: str, **kwargs: str) -> str:
    install_args = []
    for k, v in kwargs.items():
        k = k.replace("_", "-")
        if v != "":
            install_args.append(f"    --{k}={v} \\")
        else:
            install_args.append(f"    --{k} \\")
    install_args_str = "\n".join(install_args)

    return rf"""
# Install requirements for the toolchain.
WORKDIR /opt/cp2k-toolchain
COPY ./tools/toolchain/install_requirements*.sh ./
RUN ./install_requirements.sh {base_image}

# Install the toolchain.
RUN mkdir scripts
COPY ./tools/toolchain/scripts/VERSION \
     ./tools/toolchain/scripts/tool_kit.sh \
     ./tools/toolchain/scripts/common_vars.sh \
     ./tools/toolchain/scripts/signal_trap.sh \
     ./tools/toolchain/scripts/get_openblas_arch.sh \
     ./scripts/
COPY ./tools/toolchain/install_cp2k_toolchain.sh .
RUN ./install_cp2k_toolchain.sh \
{install_args_str}
    --dry-run

# Dry-run leaves behind config files for the followup install scripts.
# This breaks up the lengthy installation into smaller build steps.
COPY ./tools/toolchain/scripts/stage0/ ./scripts/stage0/
RUN  ./scripts/stage0/install_stage0.sh && rm -rf ./build

COPY ./tools/toolchain/scripts/stage1/ ./scripts/stage1/
RUN  ./scripts/stage1/install_stage1.sh && rm -rf ./build

COPY ./tools/toolchain/scripts/stage2/ ./scripts/stage2/
RUN  ./scripts/stage2/install_stage2.sh && rm -rf ./build

COPY ./tools/toolchain/scripts/stage3/ ./scripts/stage3/
RUN  ./scripts/stage3/install_stage3.sh && rm -rf ./build

COPY ./tools/toolchain/scripts/stage4/ ./scripts/stage4/
RUN  ./scripts/stage4/install_stage4.sh && rm -rf ./build

COPY ./tools/toolchain/scripts/stage5/ ./scripts/stage5/
RUN  ./scripts/stage5/install_stage5.sh && rm -rf ./build

COPY ./tools/toolchain/scripts/stage6/ ./scripts/stage6/
RUN  ./scripts/stage6/install_stage6.sh && rm -rf ./build

COPY ./tools/toolchain/scripts/stage7/ ./scripts/stage7/
RUN  ./scripts/stage7/install_stage7.sh && rm -rf ./build

COPY ./tools/toolchain/scripts/stage8/ ./scripts/stage8/
RUN  ./scripts/stage8/install_stage8.sh && rm -rf ./build

COPY ./tools/toolchain/scripts/stage9/ ./scripts/stage9/
RUN  ./scripts/stage9/install_stage9.sh && rm -rf ./build
""".lstrip()


# ======================================================================================
def install_cp2k_spack(
    version: str,
    mpi_mode: str,
    base_image: str = "ubuntu:26.04",
    gcc_version: int | None = None,
    gpu_model: str = "none",
    feature_flags: str = "",
    testopts: str = "",
    image_tag: str = "",
) -> str:
    if gcc_version is None or "fedora" in base_image:
        gcc_compilers = "g++ gcc gfortran"
    elif "opensuse/leap" in base_image:
        gcc_compilers = f"gcc gcc{gcc_version} gcc-c++ gcc{gcc_version}-c++ gcc-fortran gcc{gcc_version}-fortran"
    elif "rockylinux" in base_image:
        gcc_compilers = f"gcc gcc-c++ gcc-fortran"
    else:
        gcc_compilers = f"g++ g++-{gcc_version} gcc gcc-{gcc_version} gfortran gfortran-{gcc_version}"
    # Use the system GCC when no version is specified.
    gcc_version_flag = "" if gcc_version is None else f"-gv {gcc_version}"
    # Use external packages if possible
    use_externals = "-ue"
    # Static CP2K builds use the GCC compiler built with spack
    if version.endswith("-static"):
        use_externals = ""
        # A spack build of the same GCC version as the installed one
        # of the host system and ignoring all externals at the same
        # time is not supported
        if gcc_version == 13:
            print(
                f"\nERROR: GCC 13 is the default version of Ubuntu 24.04 and a spack build of the same version is not possible"
            )
        gcc_compilers = f"g++ gcc gfortran"
    if mpi_mode == "openmpi":
        use_externals = ""
    # Assemble docker file
    output = (
        install_base_image(
            base_image=rf"{base_image}", gcc_compilers=gcc_compilers, stage="build"
        )
        + rf"""
ARG IMAGE_TAG
ENV IMAGE_TAG=${{IMAGE_TAG:-{image_tag}}}

ARG SPACK_CACHE
ENV SPACK_CACHE="${{SPACK_CACHE:-s3://spack-cache --s3-endpoint-url=http://host.containers.internal:9000}}"

# Build CP2K dependencies
WORKDIR /opt/cp2k
COPY ./tools/spack ./tools/spack
COPY ./tools/docker ./tools/docker
COPY ./make_cp2k.sh .
RUN ./make_cp2k.sh -bd_only -cv {version} -gpu {gpu_model} {gcc_version_flag} -mpi {mpi_mode} {use_externals} {feature_flags}

###### Stage 2: Build CP2K ######

FROM build_deps AS build_cp2k

COPY ./src ./src
COPY ./data ./data
COPY ./tools/build_utils ./tools/build_utils
COPY ./cmake ./cmake
COPY ./CMakeLists.txt .

RUN ./make_cp2k.sh -cv {version} {gcc_version_flag} -gpu {gpu_model} -mpi {mpi_mode} {feature_flags}
"""
    )
    output += (
        install_base_image(
            base_image=rf"{base_image}", gcc_compilers=gcc_compilers, stage="install"
        )
        + rf"""
WORKDIR /opt/cp2k

# Install CP2K dependencies built with spack
COPY --from=build_cp2k /opt/cp2k/spack/spack/opt/spack ./spack/spack/opt/spack

# Install CP2K
COPY --from=build_cp2k /opt/cp2k/install ./install

# Install CP2K regression tests
COPY ./tests ./tests
COPY --from=build_cp2k /opt/cp2k/src/grid/sample_tasks ./src/grid/sample_tasks

# Install CP2K/Quickstep CI benchmarks
COPY ./benchmarks/CI ./benchmarks/CI

# Do not rely only on LD_LIBRARY_PATH because it is fragile
COPY --from=build_cp2k /etc/ld.so.conf.d/cp2k.conf /etc/ld.so.conf.d/cp2k.conf
RUN ldconfig

# Run CP2K regression test
RUN /opt/cp2k/install/bin/launch /opt/cp2k/install/bin/run_tests {testopts} || echo "ERROR: Tests failed"

# Create entrypoint and finalise container build
WORKDIR /mnt
ENTRYPOINT ["/opt/cp2k/install/bin/launch"]
CMD ["cp2k", "--help", "--version"]
"""
    )
    return output


# ======================================================================================
def install_base_image(
    base_image: str,
    gcc_compilers: str,
    stage: str,
) -> str:
    if stage == "build":
        output = rf"""
ARG BASE_IMAGE="{base_image}"

###### Stage 1: Build CP2K dependencies ######

FROM "${{BASE_IMAGE}}" AS build_deps
"""
        if "fedora" in base_image:
            output += rf"""
RUN dnf -qy install \
    cmake \
    {gcc_compilers} \
    git \
    libffi-devel \
    liblsan \
    libtool \
    make \
    patch \
    perl-core \
    pkg-config \
    python3 \
    python3-devel \
    unzip \
    wget \
    zlib-devel \
    zlib-static \
    && dnf clean -q all
"""
        elif "opensuse/leap" in base_image:
            output += rf"""
RUN zypper --non-interactive --quiet ref && \
    zypper --non-interactive --quiet in --no-recommends \
    bzip2 \
    cmake \
    diffutils \
    findutils \
    {gcc_compilers} \
    git \
    gzip \
    libssh-devel \
    libopenssl-devel \
    libtool \
    lsb-release \
    make \
    patch \
    pkgconf \
    python313 \
    python313-devel \
    unzip \
    wget \
    xz \
    zstd \
    && zypper --non-interactive --quiet clean --all
"""
        elif "rockylinux" in base_image:
            output += rf"""
RUN dnf -y install dnf-plugins-core && \
    dnf --enablerepo=crb -qy install \
    bzip2 \
    cmake \
    {gcc_compilers} \
    git \
    libtool \
    openssl-devel \
    patch \
    python3 \
    python3-devel \
    python3-pip \
    unzip \
    wget \
    xz \
    zlib-devel \
    zlib-static \
    && dnf clean -q all
"""
        elif "ubuntu" in base_image:
            output += rf"""
RUN apt-get update -qq && apt-get install -qq --no-install-recommends \
    bzip2 \
    ca-certificates \
    cmake \
    {gcc_compilers} \
    git \
    gnupg \
    libssh-dev \
    libssl-dev \
    libtool \
    libtool-bin \
    lsb-release \
    make \
    patch \
    pkgconf \
    python3 \
    python3-dev \
    python3-pip \
    python3-venv \
    unzip \
    wget \
    xxd \
    xz-utils \
    zlib1g \
    zlib1g-dev \
    zstd \
    && rm -rf /var/lib/apt/lists/*
"""
        else:
            print(f"\nERROR: Unknown base image {base_image} specified\n")
        if "nvidia" in base_image:
            output += rf"""
# Setup CUDA environment
ENV LD_LIBRARY_PATH /usr/local/cuda/lib64

# Disable JIT cache as there seems to be an issue with file locking on overlayfs
# See also https://github.com/cp2k/cp2k/pull/2337
ENV CUDA_CACHE_DISABLE 1
"""
    elif stage == "install":
        output = rf"""
###### Stage 3: Install CP2K ######

FROM "${{BASE_IMAGE}}" AS install_cp2k
"""
        if "fedora" in base_image:
            output += rf"""
RUN dnf -qy install \
    {gcc_compilers} \
    python3 \
    liblsan \
    && dnf clean -q all
"""
        elif "opensuse/leap" in base_image:
            output += rf"""
RUN zypper --non-interactive --quiet ref && \
    zypper --non-interactive --quiet in --no-recommends \
    {gcc_compilers} \
    python313 \
    && zypper --non-interactive --quiet clean --all
"""
        elif "rockylinux" in base_image:
            output += rf"""
RUN dnf -y install dnf-plugins-core && \
    dnf --enablerepo=crb -qy install \
    {gcc_compilers} \
    python3 \
    python3-pip \
    && dnf clean -q all
"""
        elif "ubuntu" in base_image:
            output += rf"""
RUN apt-get update -qq && apt-get install -qq --no-install-recommends \
    {gcc_compilers} \
    python3 \
    && rm -rf /var/lib/apt/lists/*
"""
        else:
            print(f"\nERROR: Unknown base image {base_image} specified\n")
        if "nvidia" in base_image:
            output += rf"""
# Setup CUDA environment
ENV LD_LIBRARY_PATH /usr/local/cuda/lib64

# Disable JIT cache as there seems to be an issue with file locking on overlayfs
# See also https://github.com/cp2k/cp2k/pull/2337
ENV CUDA_CACHE_DISABLE 1
"""
    else:
        print(f"\nERROR: Unknown stage {stage} specified\n")
    return output


# ======================================================================================
class OutputFile:
    def __init__(self, filename: str, check: bool) -> None:
        self.filename = filename
        self.check = check
        self.content = io.StringIO()
        self.content.write(f"#\n")
        self.content.write(f"# This file was created by generate_dockerfiles.py.\n")
        if "_spack_" in filename:
            self.image_tag = filename.removeprefix("Dockerfile.test_")
            usage = f"./spack_cache_start.sh; podman build --shm-size=1g -t {self.image_tag.lower()} -f ./{filename} ../../"
        else:
            self.image_tag = ""
            usage = f"podman build --shm-size=1g -f ./{filename} ../../"
        self.content.write(f"# Usage: {usage}\n#\n")

    def __enter__(self) -> "OutputFile":
        return self

    def write(self, text: str) -> None:
        self.content.write(text)

    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
        output_path = Path(__file__).parent / self.filename
        if self.check:
            assert output_path.read_text(encoding="utf8") == self.content.getvalue()
            print(f"File {output_path} is consistent with generator script.")
        else:
            output_path.write_text(self.content.getvalue(), encoding="utf8")
            print(f"Wrote {output_path}")


# ======================================================================================
main()

# EOF