FROM nvidia/cuda:12.8.0-devel-ubuntu22.04

# Set environment variables
ENV DEBIAN_FRONTEND=noninteractive
ENV HOME=/root
ENV CUDA_HOME=/usr/local/cuda
ENV PATH=${CUDA_HOME}/bin:${PATH}
ENV LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${LD_LIBRARY_PATH}
ENV TORCH_CUDA_ARCH_LIST="12.0"

# Install system dependencies and build tools
RUN apt-get update && \
    apt-get install -y \
    curl wget gpg apt-transport-https git software-properties-common \
    build-essential cmake ninja-build \
    libopenblas-dev libomp-dev \
    && rm -rf /var/lib/apt/lists/*

# Install Python 3.12 (as recommended)
RUN add-apt-repository ppa:deadsnakes/ppa && \
    apt-get update && \
    apt-get install -y python3.11 python3.11-venv python3.11-dev && \
    update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.11 1 && \
    update-alternatives --set python3 /usr/bin/python3.11 && \
    curl -sS https://bootstrap.pypa.io/get-pip.py | python3.11

# Install uv for faster dependency management
RUN curl -LsSf https://astral.sh/uv/install.sh | sh && \
    echo 'export PATH="$HOME/.local/bin:$PATH"' >> ~/.bashrc

# Create workspace
RUN mkdir -p /workspace
WORKDIR /workspace

# CRITICAL: Install in the EXACT order specified in the README
# Step 1: Create venv
RUN python3 -m venv /opt/unsloth-env
ENV PATH="/opt/unsloth-env/bin:$PATH"

# Step 2: Install vllm with cu128 (MUST be first)
RUN /opt/unsloth-env/bin/pip install --upgrade pip setuptools wheel && \
    /opt/unsloth-env/bin/pip install -U vllm --extra-index-url https://wheels.vllm.ai/nightly

# Step 3: Install unsloth dependencies
RUN /opt/unsloth-env/bin/pip install unsloth unsloth_zoo bitsandbytes

# Step 4: Build xformers from source (no Blackwell wheels exist yet)
RUN git clone --depth=1 https://github.com/facebookresearch/xformers --recursive /tmp/xformers && \
    cd /tmp/xformers && \
    /opt/unsloth-env/bin/pip uninstall -y xformers && \
    /opt/unsloth-env/bin/python setup.py install && \
    rm -rf /tmp/xformers

# Step 5: Update triton to >=3.3.1 for Blackwell
RUN /opt/unsloth-env/bin/pip install -U "triton>=3.3.1"

# Step 6: Pin transformers to avoid gradient checkpointing bug
RUN /opt/unsloth-env/bin/pip install -U "transformers==4.52.4"

# Step 7: Might need to downgrade numpy
RUN /opt/unsloth-env/bin/pip install "numpy<=2.2"

# Now install your dev tools (code-server, Node.js, etc)
RUN curl -fsSL https://deb.nodesource.com/setup_22.x | bash - && \
    apt-get install -y nodejs && \
    curl -fsSL https://code-server.dev/install.sh | sh && \
    npm install -g @anthropic-ai/claude-code @anthropic-ai/dxt && \
    apt-get clean && \
    rm -rf /var/lib/apt/lists/*

# Install Ollama
RUN curl -fsSL https://ollama.com/install.sh | sh || true

# Configure code-server
RUN mkdir -p /root/.config/code-server /root/.ollama && \
    echo "bind-addr: 0.0.0.0:8443\nauth: none\ncert: false" > /root/.config/code-server/config.yaml

# Install VS Code extensions
RUN code-server --install-extension ms-python.python && \
    code-server --install-extension ritwickdey.LiveServer && \
    code-server --install-extension ms-toolsai.jupyter

# Create startup script with proper env activation
RUN echo '#!/bin/bash\n\
# Activate the unsloth environment\n\
source /opt/unsloth-env/bin/activate\n\
\n\
# Start Ollama in the background\n\
/usr/local/bin/ollama serve &\n\
\n\
# Give Ollama a moment to start\n\
sleep 2\n\
\n\
# Start code-server with the activated environment\n\
exec code-server --disable-telemetry --bind-addr 0.0.0.0:8443 /workspace\n\
' > /start.sh && \
    chmod +x /start.sh

# Expose ports
EXPOSE 8443 11434

# Copy requirements.txt to workspace (will fail build if not present)
COPY requirements.txt /workspace/requirements.txt

# Create a script to install from requirements.txt with proper environment
RUN echo '#!/bin/bash\n\
echo "Installing from requirements.txt..."\n\
source /opt/unsloth-env/bin/activate\n\
pip install -r /workspace/requirements.txt\n\
echo "Installation complete!"\n\
' > /workspace/install-requirements.sh && \
    chmod +x /workspace/install-requirements.sh

# Optional: Install requirements.txt if you want it done at build time
# Uncomment the next line if you want automatic installation
RUN /opt/unsloth-env/bin/pip install -r /workspace/requirements.txt

CMD ["/start.sh"]