# Use Ubuntu + CUDA base image for full GPU support
FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04

# Set environment variables
ENV DEBIAN_FRONTEND=noninteractive
ENV PATH=/root/miniconda3/bin:$PATH
ENV CONDA_ENV_NAME=cosmos-predict1

# Install dependencies
RUN apt-get update && apt-get install -y \
    git curl wget sudo build-essential ca-certificates libglib2.0-0 libsm6 libxext6 libxrender-dev \
    && rm -rf /var/lib/apt/lists/*

# Install Miniconda
RUN wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O miniconda.sh && \
    bash miniconda.sh -b -p /root/miniconda3 && \
    rm miniconda.sh && \
    conda clean -ya

# Copy project files
WORKDIR /workspace
COPY . .

# Create and activate conda environment
RUN conda env create -f cosmos-predict1.yaml && \
    echo "conda activate cosmos-predict1" >> ~/.bashrc

# Activate env and install Python packages
SHELL ["conda", "run", "-n", "cosmos-predict1", "/bin/bash", "-c"]

# Install remaining pip requirements
RUN pip install -r requirements.txt

# Patch Transformer Engine headers
RUN ln -sf $CONDA_PREFIX/lib/python3.10/site-packages/nvidia/*/include/* $CONDA_PREFIX/include/ && \
    ln -sf $CONDA_PREFIX/lib/python3.10/site-packages/nvidia/*/include/* $CONDA_PREFIX/include/python3.10

# Install Transformer Engine
RUN pip install transformer-engine[pytorch]==1.12.0

# Install Apex
RUN git clone https://github.com/NVIDIA/apex && \
    CUDA_HOME=$CONDA_PREFIX pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation \
        --config-settings "--build-option=--cpp_ext" \
        --config-settings "--build-option=--cuda_ext" \
        ./apex

# Install MoGe
RUN pip install git+https://github.com/microsoft/MoGe.git

# (Optional) Test the environment
RUN CUDA_HOME=$CONDA_PREFIX PYTHONPATH=$(pwd) python scripts/test_environment.py

# Set the command to run app (replace with your app script or gradio interface)
CMD ["conda", "run", "--no-capture-output", "-n", "cosmos-predict1", "python", "app.py"]