# Builder stage
FROM nvidia/cuda:12.1.0-cudnn8-runtime-ubuntu22.04

# Add these lines after the FROM statement

# System-level configuration
ENV DEBIAN_FRONTEND=noninteractive \
    PYTHONUNBUFFERED=1 \
    CUDA_HOME=/usr/local/cuda \
    PATH=/usr/local/cuda/bin:$PATH \
    LD_LIBRARY_PATH=/usr/local/cuda/lib64:/usr/local/cuda/extras/CUPTI/lib64:$LD_LIBRARY_PATH \
    TORCH_CUDA_ARCH_LIST="8.6" \
    PIP_NO_CACHE_DIR=1 \
    TORCH_NVCC_FLAGS="-Xfatbin -compress-all" \
    MAX_JOBS=4

# Install build dependencies
RUN apt-get update && apt-get install -y --no-install-recommends \
    python3-pip \
    && apt-get clean \
    && rm -rf /var/lib/apt/lists/* \
    && python3 -m pip install --no-cache-dir --upgrade pip


# Set up Python
RUN ln -sf /usr/bin/python3 /usr/bin/python && \
    curl -sS https://bootstrap.pypa.io/get-pip.py | python3 && \
    python3 -m pip install --no-cache-dir pip==23.3.1 setuptools==69.0.3 wheel==0.42.0 && \
    rm -rf /root/.cache/pip/*

# Replace the pip install command (around line 80) with:
RUN python3 -m pip install --no-cache-dir -r requirements.txt && \
    rm -rf /root/.cache/pip/* && \
    find /usr/local/lib/python3.* -name '*.pyc' -delete && \
    find /usr/local/lib/python3.* -name '__pycache__' -exec rm -r {} +

# Runtime stage
FROM nvidia/cuda:12.1.0-cudnn8-runtime-ubuntu22.04

# Copy Python environment from builder
COPY --from=builder /usr/local /usr/local
COPY --from=builder /usr/bin/python* /usr/bin/
COPY --from=builder /usr/lib/python* /usr/lib/
COPY --from=builder /usr/bin/ffmpeg /usr/bin/
COPY --from=builder /usr/lib/x86_64-linux-gnu/libsndfile* /usr/lib/x86_64-linux-gnu/

# Runtime environment configuration
ENV PYTHONUNBUFFERED=1 \
    CUDA_HOME=/usr/local/cuda \
    PATH=/usr/local/cuda/bin:$PATH \
    LD_LIBRARY_PATH=/usr/local/cuda/lib64:/usr/local/cuda/extras/CUPTI/lib64:$LD_LIBRARY_PATH \
    TORCH_CUDA_ARCH_LIST="8.6" \
    NVIDIA_VISIBLE_DEVICES=0 \
    NVIDIA_DRIVER_CAPABILITIES=compute,utility,video \
    CUDA_MODULE_LOADING=LAZY \
    PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:1024 \
    OMP_NUM_THREADS=8 \
    MKL_NUM_THREADS=8 \
    CUDA_LAUNCH_BLOCKING=0 \
    NCCL_P2P_DISABLE=1

# Set up application directory
WORKDIR /app
RUN mkdir -p /app/cache /app/src && \
    chmod 777 /app/cache && \
    rm -rf /tmp/* /var/tmp/*

# Install dependencies
COPY requirements.txt /app/src/
WORKDIR /app/src
RUN pip3 install --no-cache-dir -r requirements.txt && \
    rm -rf /root/.cache/pip/* && \
    find /usr/local/lib/python3.* -name '*.pyc' -delete && \
    find /usr/local/lib/python3.* -name '__pycache__' -exec rm -r {} +

# Copy application code
COPY . /app/src/

# Remove unnecessary files
RUN rm -rf \
    /app/src/.git* \
    /app/src/*.pyc \
    /app/src/__pycache__ \
    /app/src/*.log \
    /app/src/*.tmp \
    /root/.cache/* \
    /tmp/* \
    /var/tmp/* \
    /var/cache/apt/* \
    /var/lib/apt/lists/* \
    /usr/share/doc/* \
    /usr/share/man/* \
    /usr/local/share/doc/* \
    /usr/local/share/man/*

# Runtime environment
ENV MODEL_PATH=/app/cache \
    PYTHONPATH=/app/src:$PYTHONPATH

# GPU verification
RUN python3 -c '\
import torch; \
import os; \
assert torch.cuda.is_available(), "CUDA unavailable"; \
print(f"PyTorch version: {torch.__version__}"); \
print(f"CUDA version: {torch.version.cuda}"); \
print(f"CUDA device count: {torch.cuda.device_count()}"); \
device = torch.cuda.current_device(); \
print(f"Current device: {torch.cuda.get_device_name(device)}"); \
print(f"Device memory: {torch.cuda.get_device_properties(device).total_memory / 1024**3:.2f} GB"); \
print(f"Architecture: {torch.cuda.get_device_capability(device)}");'

EXPOSE 8000

# Add healthcheck
HEALTHCHECK --interval=30s --timeout=10s --retries=3 \
    CMD curl -f http://localhost:8000/health || exit 1

CMD ["python", "server.py"]