#!/usr/bin/env bash

# Exit immediately if any command fails.
set -e
set -o pipefail


if [ $# -ne 2 ]; then
    echo "Usage: $0 <CONDA_VERSION> linux-64|osx-64|osx-arm64"
    exit 1
fi

# --- CONFIGURATIONS ---
# Make sure to update these values as appropriate
CONDA_VERSION="$1"

if [ "$2" == "linux-64" ]; then
    BUILD_DIRS="/tmp /data /mnt"
    INSTALLER="Miniforge3-Linux-x86_64.sh"
elif [ "$2" == "osx-64" ]; then
    BUILD_DIRS="/tmp "
    INSTALLER="Miniforge3-MacOSX-x86_64.sh"
elif [ "$2" == "osx-arm64" ]; then
    BUILD_DIRS="/tmp "
    INSTALLER="Miniforge3-MacOSX-arm64.sh"
else
    echo "Second argument must be one of: osx-arm64, osx-64, linux-64"
    exit 1
fi


S3_UPLOAD_PATH="<s3://XXXXXX/metaflow/condav2/conda-local/>"
S3_REMOTE_UPLOAD_PATH="<s3://XXXXXX/metaflow/condav2/conda-remote/>"


# Download the installer if it doesn't exist
if [ ! -f "$INSTALLER" ]; then
    wget "<https://github.com/conda-forge/miniforge/releases/latest/download/${INSTALLER}>"
fi

for dir in $BUILD_DIRS; do
    BUILD_DIR="$dir"

    BUILD_SUFFIX=$(basename "${BUILD_DIR}") # Extracts 'tmp' from '/tmp'
    CONDA_LOCAL_PATH="${BUILD_DIR}/metaflow-condav2-${CONDA_VERSION}"
    CONDA_LOCAL_DIST="conda-${2}-${BUILD_SUFFIX}-${CONDA_VERSION}.tgz" # Changed extension to .tgz

    echo "Starting Conda Distribution Build for version ${CONDA_VERSION}..."
    echo "Build directory set to: ${BUILD_DIR}"
    echo "Final artifact will be named: ${CONDA_LOCAL_DIST}"

    # Create the target directory and ensure we have permissions
    # mkdir -p "$(dirname "$CONDA_LOCAL_PATH")"
    # chown -R "$(whoami)" "$(dirname "$CONDA_LOCAL_PATH")"

    # Run the installer and clean up
    bash "$INSTALLER" -b -p "$CONDA_LOCAL_PATH"

    # Create the .condarc configuration file
    cat <<EOF > "${CONDA_LOCAL_PATH}/.condarc"
channels:
- conda-forge
pip_interop_enabled: true
remote_max_retries: 10
channel_priority: strict
repodata_use_zst: true
EOF

    # Copy the configuration for Mamba
    cp "${CONDA_LOCAL_PATH}/.condarc" "${CONDA_LOCAL_PATH}/.mambarc"

    # Install and update core packages
    MAMBA_EXE="${CONDA_LOCAL_PATH}/bin/mamba"

    "$MAMBA_EXE" -r "$CONDA_LOCAL_PATH" update -n base mamba conda pip -y
    "$MAMBA_EXE" -r "$CONDA_LOCAL_PATH" install -n base \
    "conda-lock==3.0.4" \
    "conda-package-handling" \
    "micromamba>=2.3.2" \
    "boto3" \
    "uv" \
    -y

    PYTHON_VERSION=$($CONDA_LOCAL_PATH/bin/python -c "import sys; print(f'python{sys.version_info.major}.{sys.version_info.minor}')")

    # Apply conda-lock patch (remove when no longer needed) and update conda-lock
    # version to be unpinned.    
    SITE_PACKAGES_DIR="${CONDA_LOCAL_PATH}/lib/${PYTHON_VERSION}/site-packages"
    PATCH_FILE="$(cd "$(dirname "${BASH_SOURCE[0]}")/../operations" && pwd)/conda-lock_v3.0.4.diff"
    ORIG_DIR="$(pwd)"
    cd "${SITE_PACKAGES_DIR}" || { echo "Failed to cd to ${SITE_PACKAGES_DIR}"; exit 1; }
    patch -p1 < "${PATCH_FILE}"
    cd "${ORIG_DIR}"
    
    # Isolate the environment by modifying Conda's search path
    CONSTANTS_PY_PATH="${CONDA_LOCAL_PATH}/lib/${PYTHON_VERSION}/site-packages/conda/base/constants.py"

    sed -i.bak "/DEFAULT_CHANNEL_ALIAS/i\\
    SEARCH_PATH = ('\$CONDA_ROOT/.condarc',)
    " "$CONSTANTS_PY_PATH"

    # Clean up caches to reduce size
    "$MAMBA_EXE" -r "$CONDA_LOCAL_PATH" clean --all -y

    # Create the final tarball
    cd "$CONDA_LOCAL_PATH"

    if [[ "$(uname)" == "Darwin" ]]; then
        tar --disable-copyfile -czf "${BUILD_DIR}/${CONDA_LOCAL_DIST}" .
    else
        tar -czf "${BUILD_DIR}/${CONDA_LOCAL_DIST}" .
    fi

    cd -

    echo "✅ Build complete!"
    echo "Distribution created at: ${BUILD_DIR}/${CONDA_LOCAL_DIST}"

    # Upload the final artifact to S3
    echo "Uploading ${CONDA_LOCAL_DIST} to ${S3_UPLOAD_PATH}..."
    aws s3 cp "${BUILD_DIR}/${CONDA_LOCAL_DIST}" "${S3_UPLOAD_PATH}${CONDA_LOCAL_DIST}"

    # Upload the micromamba binary as the remote installer
    echo "Uploading micromamba binary to ${S3_REMOTE_UPLOAD_PATH}..."
    aws s3 cp "${BUILD_DIR}/metaflow-condav2-${CONDA_VERSION}/bin/micromamba" "${S3_REMOTE_UPLOAD_PATH}conda-${2}-${CONDA_VERSION}"

    echo "🚀 Upload complete!"
done

rm "$INSTALLER"

CONDA_REMOTE_INSTALLER_DIRNAME = from_conf(
    "CONDA_REMOTE_INSTALLER_DIRNAME", "conda-remote"
)

# Override when using a consistent naming scheme for installed conda distribution.
# Note: set with METAFLOW__CONDA_VERSION_ID (double __)
_CONDA_VERSION_ID = from_conf("_CONDA_VERSION_ID", "20251022")

# Ditto for local installation of conda. If not specified,
# a conda executable needs to be installed and will be used (if locatable in PATH)
CONDA_LOCAL_DIST_DIRNAME = from_conf("CONDA_LOCAL_DIST_DIRNAME", "conda-local")

# Binary within CONDA_REMOTE_INSTALLER_DIRNAME to use as the binary on remote instances.
# Use {arch} to specify the architecture. This should be fully functional binary
# To regenerate this binary, download micromamba for the proper platform and upload
# it here
CONDA_REMOTE_INSTALLER = from_conf(
    "CONDA_REMOTE_INSTALLER", f"conda-{{arch}}-{_CONDA_VERSION_ID}"
)

CONDA_LOCAL_DIST = from_conf(
    "CONDA_LOCAL_DIST", f"conda-{{arch}}-{_conda_dir_name()}-{_CONDA_VERSION_ID}.tgz"
)

CONDA_LOCAL_PATH = from_conf(
    "CONDA_LOCAL_PATH",
    os.path.join("/", _conda_dir_name(), f"metaflow-condav2-{_CONDA_VERSION_ID}"),
)