# Dockerfile.gpu - CUDA-enabled PaddleOCR REST API
#
# Supports both architectures:
# - x86_64: Uses paddlepaddle-gpu from PaddlePaddle's CUDA index
# - ARM64: Uses local wheel from ./wheels/ (built on DGX Spark)
#
# For ARM64 (DGX Spark), first build the wheel:
#   docker compose --profile build run --rm build-paddle
# Then build this image:
#   docker compose build ocr-gpu
#
# For x86_64, just build directly (no wheel needed):
#   docker compose build ocr-gpu

FROM nvidia/cuda:12.4.1-cudnn-runtime-ubuntu22.04

LABEL maintainer="Sergio Jimenez"
LABEL description="PaddleOCR Tuning REST API - GPU/CUDA version"

WORKDIR /app

# Set environment variables
ENV DEBIAN_FRONTEND=noninteractive
ENV PYTHONUNBUFFERED=1
ENV CUDA_VISIBLE_DEVICES=0

# 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 \
    libgl1 \
    libglib2.0-0 \
    libsm6 \
    libxext6 \
    libxrender1 \
    libgomp1 \
    && rm -rf /var/lib/apt/lists/* \
    && ln -sf /usr/bin/python3.11 /usr/bin/python

# Copy local wheels directory (may contain ARM64 wheel from build-paddle)
COPY wheels/ /tmp/wheels/

# Copy requirements
COPY requirements-gpu.txt .

# Install paddlepaddle: prefer local wheel (ARM64), fallback to CUDA index (x86_64)
RUN if ls /tmp/wheels/paddlepaddle*.whl 1>/dev/null 2>&1; then \
        echo "=== Installing PaddlePaddle from local wheel (ARM64) ===" && \
        pip install --no-cache-dir /tmp/wheels/paddlepaddle*.whl; \
    else \
        echo "=== Installing PaddlePaddle from CUDA index (x86_64) ===" && \
        pip install --no-cache-dir paddlepaddle-gpu==3.2.0 -i https://www.paddlepaddle.org.cn/packages/stable/cu126/; \
    fi

# Install remaining dependencies (skip paddlepaddle-gpu line from requirements)
RUN grep -v "paddlepaddle-gpu" requirements-gpu.txt > /tmp/requirements-no-paddle.txt && \
    pip install --no-cache-dir -r /tmp/requirements-no-paddle.txt && \
    rm -rf /tmp/wheels /tmp/requirements-no-paddle.txt

# Copy application code
COPY paddle_ocr_tuning_rest.py .
COPY dataset_manager.py .

# Build arguments for models to bake into image
ARG DET_MODEL=PP-OCRv5_server_det
ARG REC_MODEL=PP-OCRv5_server_rec

# Set as environment variables (can be overridden at runtime)
ENV PADDLE_DET_MODEL=${DET_MODEL}
ENV PADDLE_REC_MODEL=${REC_MODEL}

# Download models during build (not at runtime)
RUN python -c "\
import os; \
from paddleocr import PaddleOCR; \
det = os.environ.get('PADDLE_DET_MODEL', 'PP-OCRv5_server_det'); \
rec = os.environ.get('PADDLE_REC_MODEL', 'PP-OCRv5_server_rec'); \
print(f'Downloading models: det={det}, rec={rec}'); \
ocr = PaddleOCR(text_detection_model_name=det, text_recognition_model_name=rec); \
print('Models downloaded successfully!')"

# Volume for dataset and optional additional model cache
VOLUME ["/app/dataset", "/root/.paddlex"]

# Expose API port
EXPOSE 8000

# Health check
HEALTHCHECK --interval=30s --timeout=10s --start-period=60s --retries=3 \
    CMD python -c "import urllib.request; urllib.request.urlopen('http://localhost:8000/health')" || exit 1

# Run the API server
CMD ["uvicorn", "paddle_ocr_tuning_rest:app", "--host", "0.0.0.0", "--port", "8000"]