# Multi-stage build for optimal image size

# Stage 1: Builder
FROM nvidia/cuda:11.8.0-runtime-ubuntu22.04 as builder

WORKDIR /build

# Install system dependencies
RUN apt-get update && apt-get install -y --no-install-recommends \
    python3.10 \
    python3-pip \
    python3.10-dev \
    git \
    build-essential \
    && rm -rf /var/lib/apt/lists/*

# Copy requirements
COPY requirements.txt .

# Create wheels for all dependencies
RUN pip3 install --no-cache-dir --user --upgrade pip setuptools wheel && \
    pip3 install --no-cache-dir --user -r requirements.txt

# Stage 2: Runtime
FROM nvidia/cuda:11.8.0-runtime-ubuntu22.04

WORKDIR /app

# Install minimal runtime dependencies
RUN apt-get update && apt-get install -y --no-install-recommends \
    python3.10 \
    python3-pip \
    libsm6 \
    libxext6 \
    libxrender-dev \
    libgomp1 \
    && rm -rf /var/lib/apt/lists/*

# Copy Python packages from builder
COPY --from=builder /root/.local /root/.local

# Set Python path
ENV PATH=/root/.local/bin:$PATH \
    PYTHONUNBUFFERED=1 \
    PYTHONDONTWRITEBYTECODE=1 \
    TORCH_HOME=/app/.cache/torch

# Create cache directory
RUN mkdir -p /app/.cache/torch

# Copy application files
COPY app.py .
COPY config.json .
COPY model_weights/ ./model_weights/
ENV GENERATOR_CHECKPOINT=/app/model_weights/generator_best.pth

# Health check
HEALTHCHECK --interval=30s --timeout=10s --start-period=40s --retries=3 \
    CMD python3 -c "import torch; print('OK')" || exit 1

# Expose port
EXPOSE 7860

# Set up non-root user (optional, for security)
RUN useradd -m -u 1000 appuser && \
    chown -R appuser:appuser /app
USER appuser

# Run app
CMD ["python3", "app.py", "--server_name", "0.0.0.0", "--server_port", "7860"]