# Dockerfile.gpu - Multi-stage GPU Dockerfile
#
# Build base only (push to registry, rarely changes):
#   docker build --target base -t seryus.ddns.net/unir/paddle-ocr-gpu-base:latest -f Dockerfile.gpu .
#
# Build deploy (uses base, fast - code only):
#   docker build --target deploy -t seryus.ddns.net/unir/paddle-ocr-gpu:latest -f Dockerfile.gpu .
#
# Or build all at once:
#   docker build -t paddle-ocr-api:gpu -f Dockerfile.gpu .

# =============================================================================
# STAGE 1: BASE - All dependencies (rarely changes)
# =============================================================================
FROM nvidia/cuda:12.4.1-cudnn-runtime-ubuntu22.04 AS base

LABEL maintainer="Sergio Jimenez"
LABEL description="PaddleOCR Base Image - GPU/CUDA dependencies"

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

# Fix cuDNN library path for ARM64 only (PaddlePaddle looks in /usr/local/cuda/lib64)
# x86_64 doesn't need this - PyPI wheel handles paths correctly
RUN if [ "$(uname -m)" = "aarch64" ]; then \
        mkdir -p /usr/local/cuda/lib64 && \
        ln -sf /usr/lib/aarch64-linux-gnu/libcudnn*.so* /usr/local/cuda/lib64/ && \
        ln -sf /usr/lib/aarch64-linux-gnu/libcudnn.so.9 /usr/local/cuda/lib64/libcudnn.so && \
        ldconfig; \
    fi

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

# 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) ===" && \
        python -m pip install --no-cache-dir /tmp/wheels/paddlepaddle*.whl; \
    else \
        echo "=== Installing PaddlePaddle from CUDA index (x86_64) ===" && \
        python -m pip install --no-cache-dir paddlepaddle-gpu==3.2.0 -i https://www.paddlepaddle.org.cn/packages/stable/cu126/; \
    fi && \
    rm -rf /tmp/wheels

# Install remaining dependencies
RUN python -m pip install --no-cache-dir \
    paddleocr==3.3.2 \
    jiwer \
    numpy \
    fastapi \
    "uvicorn[standard]" \
    pydantic \
    Pillow

# =============================================================================
# STAGE 2: DEPLOY - Application code (changes frequently)
# =============================================================================
FROM base AS deploy

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

WORKDIR /app

# Copy application code (this is the only layer that changes frequently)
COPY paddle_ocr_tuning_rest.py .
COPY dataset_manager.py .

# Build arguments for models
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}

# Volume for dataset and 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"]
