# Use CUDA base image without CuDNN
FROM nvidia/cuda:12.3.2-runtime-ubuntu22.04

# Set environment variables
ENV DEBIAN_FRONTEND=noninteractive

# Create a user and set environment variables
RUN useradd -m -u 1000 user
USER user

ENV HOME=/home/user \
    CUDA_HOME=/usr/local/cuda \
    PATH=/home/user/.local/bin:$PATH \
    LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${LD_LIBRARY_PATH} \
    LIBRARY_PATH=${CUDA_HOME}/lib64/stubs:${LIBRARY_PATH} \
    PYTHONPATH=$HOME/app \
    PYTHONUNBUFFERED=1 \
    GRADIO_ALLOW_FLAGGING=never \
    GRADIO_NUM_PORTS=1 \
    GRADIO_SERVER_NAME=0.0.0.0 \
    GRADIO_THEME=huggingface \
    GRADIO_SHARE=False \
    SYSTEM=spaces

# Set the working directory to the user's home directory
WORKDIR $HOME/app

# Switch to root to install system dependencies
USER root

# Install basic dependencies and add deadsnakes repository for Python 3.9
RUN apt-get update && apt-get install -y \
    wget \
    curl \
    git \
    software-properties-common \
    gnupg \
    libgl1-mesa-glx \
    libglib2.0-0 \
    && add-apt-repository ppa:deadsnakes/ppa -y \
    && apt-get update && apt-get install -y \
    python3.9 \
    python3.9-distutils \
    python3.9-dev \
    && curl -sS https://bootstrap.pypa.io/get-pip.py | python3.9 \
    && update-alternatives --install /usr/bin/python python /usr/bin/python3.9 1 \
    && update-alternatives --install /usr/bin/pip pip /usr/local/bin/pip 1 \
    && apt-get clean

# Install CuDNN 9.3.0
RUN apt-get update && apt-get install -y \
    wget \
    && wget https://developer.download.nvidia.com/compute/cudnn/9.3.0/local_installers/cudnn-local-repo-ubuntu2204-9.3.0_1.0-1_amd64.deb \
    && dpkg -i cudnn-local-repo-ubuntu2204-9.3.0_1.0-1_amd64.deb \
    && cp /var/cudnn-local-repo-ubuntu2204-9.3.0/cudnn-local-D9334AA3-keyring.gpg /usr/share/keyrings/ \
    && apt-get update \
    && rm cudnn-local-repo-ubuntu2204-9.3.0_1.0-1_amd64.deb

# Check NVIDIA packages installed via pip (using || true to prevent non-zero exit)
RUN echo "Checking NVIDIA packages:" && pip list | grep nvidia || true

# Fix permissions for pip cache
RUN mkdir -p /home/user/.cache/pip && chown -R user:user /home/user/.cache

COPY . .
# Install your requirements
RUN pip install --upgrade pip
RUN pip install -r requirements.txt
RUN pip install gradio

# Verify CuDNN installation
RUN ldconfig && echo "Testing CuDNN installation..." \
    && if [ -f /usr/local/cuda/include/cudnn.h ]; then \
       echo "CuDNN found in /usr/local/cuda/include/cudnn.h"; \
       cat /usr/local/cuda/include/cudnn.h | grep CUDNN_MAJOR -A 2; \
    fi

USER user

# Set the environment variable to specify the GPU device
ENV CUDA_DEVICE_ORDER=PCI_BUS_ID
ENV CUDA_VISIBLE_DEVICES=0

# Command to run your application
CMD ["python", "app.py"]