# Base image with complete ML and hardware acceleration stack
FROM pytorch/pytorch:2.8.0-cuda12.6-cudnn9-devel

# Install build dependencies and system libraries
RUN apt-get update && apt-get install -y \
    # Build tools
    build-essential \
    cmake \
    git \
    pkg-config \
    wget \
    unzip \
    yasm \
    nasm \
    # Additional dependencies for FFmpeg/NVIDIA build
    libtool \
    libc6 \
    libc6-dev \
    libnuma1 \
    libnuma-dev \
    # System libraries
    libgl1-mesa-glx \
    libglib2.0-0 \
    libgomp1 \
    # Core media libraries (essential ones only)
    libjpeg-dev \
    libpng-dev \
    libx264-dev \
    libx265-dev \
    libvpx-dev \
    libmp3lame-dev \
    libv4l-dev \
    # TurboJPEG for fast JPEG encoding
    libturbojpeg0-dev \
    # Python development
    python3-dev \
    python3-numpy \
    && rm -rf /var/lib/apt/lists/*

# CUDA development tools already available in devel image

# Ensure CUDA paths are available
ENV PATH="/usr/local/cuda/bin:${PATH}"
ENV LD_LIBRARY_PATH="/usr/local/cuda/lib64:${LD_LIBRARY_PATH}"

# Install NVIDIA Video Codec SDK headers (official method)
RUN cd /tmp && \
    git clone https://git.videolan.org/git/ffmpeg/nv-codec-headers.git && \
    cd nv-codec-headers && \
    make install && \
    cd / && rm -rf /tmp/*

# Build FFmpeg from source with NVIDIA CUDA support
RUN cd /tmp && \
    echo "Building FFmpeg with NVIDIA CUDA support..." && \
    # Download FFmpeg source (official method)
    git clone https://git.ffmpeg.org/ffmpeg.git ffmpeg/ && \
    cd ffmpeg && \
    # Configure with NVIDIA support (following official NVIDIA documentation)
    ./configure \
        --prefix=/usr/local \
        --enable-shared \
        --disable-static \
        --enable-nonfree \
        --enable-gpl \
        --enable-cuda-nvcc \
        --enable-cuda-llvm \
        --enable-cuvid \
        --enable-nvdec \
        --enable-nvenc \
        --enable-libnpp \
        --nvcc=/usr/local/cuda/bin/nvcc \
        --extra-cflags=-I/usr/local/cuda/include \
        --extra-ldflags=-L/usr/local/cuda/lib64 \
        --extra-libs=-lcuda \
        --extra-libs=-lcudart \
        --extra-libs=-lnvcuvid \
        --extra-libs=-lnvidia-encode \
        --enable-libx264 \
        --enable-libx265 \
        --enable-libvpx \
        --enable-libmp3lame && \
    # Build and install
    make -j$(nproc) && \
    make install && \
    ldconfig && \
    # Verify CUVID decoders are available
    echo "=== Verifying FFmpeg CUVID Support ===" && \
    (ffmpeg -hide_banner -decoders 2>/dev/null | grep cuvid || echo "No CUVID decoders found") && \
    echo "=== Verifying FFmpeg NVENC Support ===" && \
    (ffmpeg -hide_banner -encoders 2>/dev/null | grep nvenc || echo "No NVENC encoders found") && \
    cd / && rm -rf /tmp/*

# Set environment variables for maximum hardware acceleration
ENV LD_LIBRARY_PATH="/usr/local/cuda/lib64:/usr/local/lib:${LD_LIBRARY_PATH}"
ENV PKG_CONFIG_PATH="/usr/local/lib/pkgconfig:${PKG_CONFIG_PATH}"
ENV PYTHONPATH="/usr/local/lib/python3.10/dist-packages:${PYTHONPATH}"

# Optimized environment variables for hardware acceleration
ENV OPENCV_FFMPEG_CAPTURE_OPTIONS="rtsp_transport;tcp|hwaccel;cuda|hwaccel_device;0|video_codec;h264_cuvid|hwaccel_output_format;cuda"
ENV OPENCV_FFMPEG_WRITER_OPTIONS="video_codec;h264_nvenc|preset;fast|tune;zerolatency|gpu;0"
ENV CUDA_VISIBLE_DEVICES=0
ENV NVIDIA_VISIBLE_DEVICES=all
ENV NVIDIA_DRIVER_CAPABILITIES=compute,video,utility

# Copy and install base requirements (exclude opencv-python since we built from source)
COPY requirements.base.txt .
RUN grep -v opencv-python requirements.base.txt > requirements.tmp && \
    mv requirements.tmp requirements.base.txt && \
    pip install --no-cache-dir -r requirements.base.txt

# Set working directory
WORKDIR /app

# Create images directory for bind mount
RUN mkdir -p /app/images && \
    chmod 755 /app/images

# This base image will be reused for all worker builds
CMD ["python3", "-m", "fastapi", "run", "--host", "0.0.0.0", "--port", "8000"]