# ============================================
# SPECMEM BRAIN CONTAINER — MINIMAL BUILD
# ============================================
# Everything SpecMem needs, nothing it doesn't:
#   - PostgreSQL 16 + pgvector
#   - Python embedding server (frankenstein-embeddings.py)
#   - Translation server (hardwick-translate.py)
#   - Mini-COT service (Pythia-410M)
#   - Supervisord (process manager)
#
# Target: ~3GB final image
# ============================================

# ============================================
# Stage 1: Python deps builder
# ============================================
FROM python:3.11-slim AS python-builder

RUN apt-get update && apt-get install -y --no-install-recommends \
    build-essential \
    && rm -rf /var/lib/apt/lists/*

# Constraints file: block ALL nvidia/CUDA packages from ever being installed.
# pip constraints are absolute — no transitive dep can override them.
RUN cat > /tmp/constraints.txt <<'EOF'
nvidia-cuda-runtime-cu12==0.0.0.invalid
nvidia-cuda-cupti-cu12==0.0.0.invalid
nvidia-cuda-nvrtc-cu12==0.0.0.invalid
nvidia-cudnn-cu12==0.0.0.invalid
nvidia-cufft-cu12==0.0.0.invalid
nvidia-curand-cu12==0.0.0.invalid
nvidia-cusolver-cu12==0.0.0.invalid
nvidia-cusparse-cu12==0.0.0.invalid
nvidia-nccl-cu12==0.0.0.invalid
nvidia-nvjitlink-cu12==0.0.0.invalid
nvidia-nvtx-cu12==0.0.0.invalid
nvidia-cuda-runtime-cu11==0.0.0.invalid
nvidia-cuda-cupti-cu11==0.0.0.invalid
nvidia-cuda-nvrtc-cu11==0.0.0.invalid
nvidia-cudnn-cu11==0.0.0.invalid
nvidia-cufft-cu11==0.0.0.invalid
nvidia-curand-cu11==0.0.0.invalid
nvidia-cusolver-cu11==0.0.0.invalid
nvidia-cusparse-cu11==0.0.0.invalid
nvidia-nccl-cu11==0.0.0.invalid
nvidia-nvjitlink-cu11==0.0.0.invalid
nvidia-nvtx-cu11==0.0.0.invalid
triton==0.0.0.invalid
EOF

# TODO: Add verification step to confirm torch CPU-only version was installed
# Step 1: Install torch CPU-only using --index-url (REPLACES PyPI as primary).
# This is the ONLY way to guarantee pip doesn't pull CUDA torch from PyPI.
# Updated to 2.4.0 for sentence-transformers compatibility
RUN pip install --no-cache-dir \
    --index-url https://download.pytorch.org/whl/cpu \
    torch==2.4.0

# Step 2: Install everything else from PyPI. torch is already satisfied
# from step 1 so pip won't re-download it. Constraints block nvidia.
RUN pip install --no-cache-dir \
    -c /tmp/constraints.txt \
    sentence-transformers \
    onnxruntime \
    optimum \
    psycopg2-binary \
    numpy \
    transformers \
    argostranslate \
    ctranslate2 \
    sentencepiece \
    emoji \
    opencc-python-reimplemented

# Verify no nvidia packages made it through
RUN pip list | grep -i nvidia && echo "NVIDIA PACKAGES FOUND - BUILD SHOULD FAIL" && exit 1 || echo "Clean: no nvidia packages"

# TODO: Consider using pip cache for faster rebuilds (--no-cache-dir currently disables it)
# Strip bloat from installed packages
# NOTE: Do NOT strip torch .so files - it breaks the _C C extension causing:
#   - NameError: name '_C' is not defined
#   - ModuleNotFoundError: No module named 'torch.testing'
RUN find /usr/local/lib/python3.11/site-packages -type d \
    \( -name "test" -o -name "tests" -o -name "__pycache__" \) \
    -exec rm -rf {} + 2>/dev/null; \
    find /usr/local/lib/python3.11/site-packages -name "*.pyc" -delete 2>/dev/null; \
    rm -rf /usr/local/lib/python3.11/site-packages/torch/utils/benchmark 2>/dev/null; \
    rm -rf /usr/local/lib/python3.11/site-packages/torch/include 2>/dev/null; \
    rm -rf /usr/local/lib/python3.11/site-packages/torch/share 2>/dev/null; \
    rm -rf /usr/local/lib/python3.11/site-packages/caffe2 2>/dev/null; \
    # Only strip .so files EXCEPT torch - torch's .so files need their symbols
    find /usr/local/lib/python3.11/site-packages -name "*.so" -not -path "*/torch/*" \
    -exec strip --strip-unneeded {} \; 2>/dev/null; \
    true

# ============================================
# Stage 2: Final minimal image
# ============================================
FROM python:3.11-slim

# Configurable container user (default: specmem)
ARG CONTAINER_USER=specmem
ENV CONTAINER_USER=${CONTAINER_USER}

ENV DEBIAN_FRONTEND=noninteractive

# PostgreSQL 16 + pgvector + supervisor — single layer, full cleanup
RUN apt-get update && apt-get install -y --no-install-recommends \
    gnupg2 lsb-release curl ca-certificates \
    && echo "deb http://apt.postgresql.org/pub/repos/apt $(lsb_release -cs)-pgdg main" \
       > /etc/apt/sources.list.d/pgdg.list \
    && curl -fsSL https://www.postgresql.org/media/keys/ACCC4CF8.asc \
       | gpg --dearmor -o /etc/apt/trusted.gpg.d/postgresql.gpg \
    && apt-get update && apt-get install -y --no-install-recommends \
    postgresql-16 \
    postgresql-16-pgvector \
    supervisor \
    && apt-get purge -y gnupg2 lsb-release curl \
    && apt-get autoremove -y --purge \
    && rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* \
    /usr/share/doc /usr/share/man /usr/share/locale \
    /var/log/* /var/cache/*

# Copy Python packages from builder (no --prefix, just site-packages)
COPY --from=python-builder /usr/local/lib/python3.11/site-packages /usr/local/lib/python3.11/site-packages
COPY --from=python-builder /usr/local/bin /usr/local/bin

WORKDIR /app

# Create non-root user BEFORE copying files so COPY --chown works
# without creating a duplicate chown layer
RUN useradd -r -m -s /bin/bash ${CONTAINER_USER} && \
    mkdir -p /data && chown ${CONTAINER_USER}:${CONTAINER_USER} /data

# Models
COPY --chown=${CONTAINER_USER}:${CONTAINER_USER} embedding-sandbox/models/ /app/models/

# Application files
COPY --chown=${CONTAINER_USER}:${CONTAINER_USER} embedding-sandbox/frankenstein-embeddings.py /app/frankenstein-embeddings.py
COPY --chown=${CONTAINER_USER}:${CONTAINER_USER} embedding-sandbox/hardwick-translate.py /app/hardwick-translate.py
COPY --chown=${CONTAINER_USER}:${CONTAINER_USER} embedding-sandbox/mini-cot-service.py /app/mini-cot-service.py

# Container config
COPY --chown=${CONTAINER_USER}:${CONTAINER_USER} container/supervisord.conf /etc/supervisor/conf.d/specmem.conf
COPY --chown=${CONTAINER_USER}:${CONTAINER_USER} container/entrypoint.sh /app/entrypoint.sh
COPY --chown=${CONTAINER_USER}:${CONTAINER_USER} container/healthcheck.sh /app/healthcheck.sh
COPY --chown=${CONTAINER_USER}:${CONTAINER_USER} container/sock_check.py /app/sock_check.py
RUN chmod +x /app/entrypoint.sh /app/healthcheck.sh

# PostgreSQL config
ENV PGDATA=${SPECMEM_CONTAINER_DATA:-/data}/pgdata
ENV PGHOST=/data/run
ENV PGPORT=5432
ENV PGUSER=${CONTAINER_USER}
ENV PGDATABASE=specmem

# Model paths
ENV SPECMEM_MODEL_PATH=/app/models/all-MiniLM-L6-v2
ENV SENTENCE_TRANSFORMERS_HOME=/app/models

# CPU-only resource limits
ENV OMP_NUM_THREADS=1
ENV MKL_NUM_THREADS=1
ENV OPENBLAS_NUM_THREADS=1

VOLUME /data
USER ${CONTAINER_USER}
ENTRYPOINT ["/app/entrypoint.sh"]