FROM mambaorg/micromamba:1.5.8

# Set environment variables
ENV DEBIAN_FRONTEND=noninteractive
ENV MAMBA_DOCKERFILE_ACTIVATE=1

# Install system dependencies as root
USER root
RUN apt-get update -y && \
    apt-get install -y --no-install-recommends \
    wget \
    git \
    procps \
    hmmer \
    curl \
    && rm -rf /var/lib/apt/lists/*

# Switch back to micromamba user
USER $MAMBA_USER

# Create conda environment with all dependencies (CPU version for compatibility)
RUN micromamba install -y -n base -c conda-forge -c bioconda -c pytorch \
    python=3.10 \
    pytorch \
    cpuonly \
    openmm \
    pdbfixer \
    anarci \
    numpy \
    "scipy>=1.6" \
    "einops>=0.3" \
    requests \
    && micromamba clean --all --yes

# Install ImmuneBuilder via pip
RUN /opt/conda/bin/pip install --no-cache-dir ImmuneBuilder

# Pre-download model weights to avoid rate limiting at runtime
USER root

# Create cache directory that ImmuneBuilder expects
RUN mkdir -p /opt/immunebuilder_weights

# Download antibody models
RUN curl -L -o /opt/immunebuilder_weights/antibody_model_1 --retry 5 --retry-delay 30 \
    "https://zenodo.org/record/7258553/files/antibody_model_1?download=1" && \
    curl -L -o /opt/immunebuilder_weights/antibody_model_2 --retry 5 --retry-delay 30 \
    "https://zenodo.org/record/7258553/files/antibody_model_2?download=1" && \
    curl -L -o /opt/immunebuilder_weights/antibody_model_3 --retry 5 --retry-delay 30 \
    "https://zenodo.org/record/7258553/files/antibody_model_3?download=1" && \
    curl -L -o /opt/immunebuilder_weights/antibody_model_4 --retry 5 --retry-delay 30 \
    "https://zenodo.org/record/7258553/files/antibody_model_4?download=1"

# Download nanobody models
RUN curl -L -o /opt/immunebuilder_weights/nanobody_model_1 --retry 5 --retry-delay 30 \
    "https://zenodo.org/record/7258553/files/nanobody_model_1?download=1" && \
    curl -L -o /opt/immunebuilder_weights/nanobody_model_2 --retry 5 --retry-delay 30 \
    "https://zenodo.org/record/7258553/files/nanobody_model_2?download=1" && \
    curl -L -o /opt/immunebuilder_weights/nanobody_model_3 --retry 5 --retry-delay 30 \
    "https://zenodo.org/record/7258553/files/nanobody_model_3?download=1" && \
    curl -L -o /opt/immunebuilder_weights/nanobody_model_4 --retry 5 --retry-delay 30 \
    "https://zenodo.org/record/7258553/files/nanobody_model_4?download=1"

# Download TCR models (TCRBuilder2+ weights)
RUN curl -L -o /opt/immunebuilder_weights/tcr_model_1 --retry 5 --retry-delay 30 \
    "https://zenodo.org/record/10581165/files/tcr_model_1?download=1" && \
    curl -L -o /opt/immunebuilder_weights/tcr_model_2 --retry 5 --retry-delay 30 \
    "https://zenodo.org/record/10581165/files/tcr_model_2?download=1" && \
    curl -L -o /opt/immunebuilder_weights/tcr_model_3 --retry 5 --retry-delay 30 \
    "https://zenodo.org/record/10581165/files/tcr_model_3?download=1" && \
    curl -L -o /opt/immunebuilder_weights/tcr_model_4 --retry 5 --retry-delay 30 \
    "https://zenodo.org/record/10581165/files/tcr_model_4?download=1"

# Set permissions and create symlinks for all users
RUN chmod -R 755 /opt/immunebuilder_weights && \
    mkdir -p /root/.cache/ImmuneBuilder && \
    ln -sf /opt/immunebuilder_weights/* /root/.cache/ImmuneBuilder/ && \
    mkdir -p /home/$MAMBA_USER/.cache/ImmuneBuilder && \
    ln -sf /opt/immunebuilder_weights/* /home/$MAMBA_USER/.cache/ImmuneBuilder/ && \
    chown -R $MAMBA_USER:$MAMBA_USER /home/$MAMBA_USER/.cache

# Create symlinks for the tools
RUN ln -sf /opt/conda/bin/ABodyBuilder2 /usr/local/bin/abodybuilder2 && \
    ln -sf /opt/conda/bin/NanoBodyBuilder2 /usr/local/bin/nanobodybuilder2 && \
    ln -sf /opt/conda/bin/TCRBuilder2 /usr/local/bin/tcrbuilder2

# Create working directory
WORKDIR /workspace

USER $MAMBA_USER

# Set default entrypoint (micromamba entrypoint activates the environment)
ENTRYPOINT ["/usr/local/bin/_entrypoint.sh"]
CMD ["ABodyBuilder2", "--help"]