# GPU-enabled worker image
# Multi-stage build: use devel image for compiling, runtime for final image

# Stage 1: Build decord with CUDA
FROM nvidia/cuda:12.1.1-cudnn8-devel-ubuntu22.04 AS builder

RUN apt-get update && apt-get install -y --no-install-recommends \
    python3.11 \
    python3.11-venv \
    python3.11-dev \
    python3-pip \
    git \
    cmake \
    build-essential \
    pkg-config \
    libavcodec-dev \
    libavformat-dev \
    libavutil-dev \
    libavdevice-dev \
    libavfilter-dev \
    libswresample-dev \
    libswscale-dev \
    && rm -rf /var/lib/apt/lists/* \
    && ln -sf /usr/bin/python3.11 /usr/bin/python3 \
    && ln -sf /usr/bin/python3 /usr/bin/python

# Download Video Codec SDK headers for NVDEC/NVCUVID
RUN git clone https://github.com/FFmpeg/nv-codec-headers.git /tmp/nv-codec-headers && \
    cd /tmp/nv-codec-headers && make install && rm -rf /tmp/nv-codec-headers

# Create stub for libnvcuvid (real library comes from driver at runtime)
RUN echo 'void* __nvcuvid_stub__;' | gcc -shared -x c - -o /usr/local/cuda/lib64/libnvcuvid.so

# Build decord with CUDA support
RUN git clone --recursive https://github.com/dmlc/decord /tmp/decord && \
    cd /tmp/decord && \
    mkdir build && cd build && \
    cmake .. -DUSE_CUDA=ON -DCMAKE_BUILD_TYPE=Release \
             -DCMAKE_CUDA_ARCHITECTURES="70;75;80;86;89;90" && \
    make -j$(nproc) && \
    cd ../python && pip install --target=/decord-install .

# Stage 2: Runtime image
FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04

WORKDIR /app

# Install Python 3.11 and system dependencies
RUN apt-get update && apt-get install -y --no-install-recommends \
    python3.11 \
    python3.11-venv \
    python3-pip \
    git \
    ffmpeg \
    && rm -rf /var/lib/apt/lists/* \
    && ln -sf /usr/bin/python3.11 /usr/bin/python3 \
    && ln -sf /usr/bin/python3 /usr/bin/python

# Upgrade pip
RUN python3 -m pip install --upgrade pip

# Install CPU dependencies first
COPY l1/requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# Install GPU-specific dependencies (CuPy for CUDA 12.x, JAX with CUDA)
RUN pip install --no-cache-dir cupy-cuda12x jax[cuda12]

# Install PyNvVideoCodec for zero-copy GPU encoding
RUN pip install --no-cache-dir PyNvVideoCodec

# Copy decord from builder stage
COPY --from=builder /decord-install /usr/local/lib/python3.11/dist-packages/
COPY --from=builder /tmp/decord/build/libdecord.so /usr/local/lib/
RUN ldconfig

# Install monorepo shared packages
COPY common/ /tmp/common/
COPY core/ /tmp/core/
RUN pip install --no-cache-dir /tmp/common/ /tmp/core/ && rm -rf /tmp/common /tmp/core

# Clone effects repo (before COPY so it gets cached)
RUN git clone https://git.rose-ash.com/art-dag/effects.git /app/artdag-effects

# Copy application (this invalidates cache for any code change)
COPY l1/ .

# Create cache directory
RUN mkdir -p /data/cache

ENV PYTHONUNBUFFERED=1
ENV PYTHONDONTWRITEBYTECODE=1
ENV EFFECTS_PATH=/app/artdag-effects
ENV PYTHONPATH=/app
# GPU persistence enabled - frames stay on GPU throughout pipeline
ENV STREAMING_GPU_PERSIST=1
# Preload libnvcuvid for decord NVDEC GPU decode
ENV LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libnvcuvid.so
# Use cluster's public IPFS gateway for HLS segment URLs
ENV IPFS_GATEWAY_URL=https://celery-artdag.rose-ash.com/ipfs

# Default command runs celery worker
CMD ["celery", "-A", "celery_app", "worker", "--loglevel=info", "-E", "-Q", "gpu,celery"]
