# thecollabagepatch/magenta:latest - Duplicable x86_64 build
FROM nvidia/cuda:12.6.2-cudnn-runtime-ubuntu22.04

# Ensure CUDA libraries are on loader path
RUN apt-get update && apt-get install -y --no-install-recommends \
    cuda-libraries-12-4 && rm -rf /var/lib/apt/lists/*
ENV LD_LIBRARY_PATH=/usr/local/cuda/lib64:/usr/local/cuda-12.4/lib64:/usr/local/cuda-12.4/compat:/usr/local/cuda/targets/x86_64-linux/lib:${LD_LIBRARY_PATH}
RUN ln -sf /usr/local/cuda/targets/x86_64-linux/lib /usr/local/cuda/lib64 || true

# Install cuDNN 9.8 for better compatibility
RUN set -eux; \
  apt-get update && apt-get install -y --no-install-recommends gnupg ca-certificates curl; \
  install -d -m 0755 /usr/share/keyrings; \
  curl -fsSL https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/3bf863cc.pub \
    | gpg --batch --yes --dearmor -o /usr/share/keyrings/cuda-archive-keyring.gpg; \
  apt-get update; \
  apt-mark unhold libcudnn9-cuda-12 || true; \
  apt-get install -y --no-install-recommends \
      'libcudnn9-cuda-12=9.8.*' \
      'libcudnn9-dev-cuda-12=9.8.*' \
      --allow-downgrades --allow-change-held-packages; \
  apt-mark hold libcudnn9-cuda-12 || true; \
  ldconfig; \
  rm -rf /var/lib/apt/lists/*

# Performance optimizations for L40S/Ada
ENV LD_PRELOAD=/usr/local/cuda/lib64/libcusparse.so.12:/usr/local/cuda/lib64/libcublas.so.12:/usr/local/cuda/lib64/libcublasLt.so.12:/usr/local/cuda/lib64/libcufft.so.11:/usr/local/cuda/lib64/libcusolver.so.11
ENV TF_GPU_ALLOCATOR=cuda_malloc_async
ENV TF_ENABLE_CUBLAS_TF32=1 NVIDIA_TF32_OVERRIDE=1

ENV DEBIAN_FRONTEND=noninteractive \
    PYTHONUNBUFFERED=1 \
    PIP_NO_CACHE_DIR=1 \
    TF_FORCE_GPU_ALLOW_GROWTH=true \
    XLA_PYTHON_CLIENT_PREALLOCATE=false \
    JAX_PLATFORMS=""

SHELL ["/bin/bash", "-c"]

# Install system dependencies
RUN apt-get update && apt-get install -y --no-install-recommends \
    software-properties-common curl ca-certificates \
    build-essential pkg-config git \
    libsndfile1 ffmpeg \
    && add-apt-repository ppa:deadsnakes/ppa -y \
    && apt-get update && apt-get install -y --no-install-recommends \
    python3.11 python3.11-dev python3.11-venv python3-pip \
    && rm -rf /var/lib/apt/lists/*

# Make python3.11 default and install uv for faster package installs
RUN ln -sf /usr/bin/python3.11 /usr/bin/python && \
    ln -sf /usr/bin/python3.11 /usr/bin/python3 && \
    python -m pip install --upgrade pip && \
    python -m pip install uv

# CRITICAL: Install TensorFlow FIRST to block tensorflow-cpu
# Using generic tf-nightly (no specific date) for duplicability
RUN uv pip install --system tf-nightly

# Install JAX with CUDA support (pinned for stability)
RUN uv pip install --system "jax[cuda12]" jaxlib

# Install base dependencies
RUN uv pip install --system \
    absl-py chex gin-config numpy requests tqdm typing-extensions \
    google-cloud-storage librosa resampy soundfile sentencepiece

# Clone and install t5x WITHOUT dependencies (avoid tensorflow-cpu)
RUN git clone https://github.com/google-research/t5x.git /t5x && \
    cd /t5x && \
    git checkout 92c5b467a5964d06c351c7eae4aa4bcd341c7ded && \
    uv pip install --system --no-deps -e .

# Install flaxformer without deps
RUN git clone https://github.com/google/flaxformer.git /flaxformer && \
    cd /flaxformer && \
    git checkout 399ea3a && \
    uv pip install --system --no-deps -e .

# Install seqio without deps and PATCH OUT tensorflow_text
RUN git clone https://github.com/google/seqio.git /seqio && \
    cd /seqio && \
    uv pip install --system --no-deps -e . && \
    # CRITICAL FIX: Remove unused tensorflow_text import (not needed by Magenta RT)
    sed -i '/import tensorflow_text as tf_text/d' /seqio/seqio/vocabularies.py

# Install airio (t5x dependency) without deps
RUN git clone https://github.com/google/airio.git /airio && \
    cd /airio && \
    uv pip install --system --no-deps -e .

# Install clu without deps
RUN git clone https://github.com/google/CommonLoopUtils.git /clu && \
    cd /clu && \
    uv pip install --system --no-deps -e .

# Now install all remaining dependencies these packages need
RUN uv pip install --system \
    flax optax orbax-checkpoint \
    fiddle cached_property tf2jax \
    aqtp etils jestimator \
    tensorflow-datasets tfds-nightly \
    apache-beam pyyaml rouge-score sacrebleu scipy \
    grain-nightly editdistance pyglove

# Patch jestimator for newer JAX - PartitionSpec moved to jax.sharding
RUN sed -i 's|from jax.experimental.pjit import PartitionSpec|from jax.sharding import PartitionSpec|g' \
    /usr/local/lib/python3.11/dist-packages/jestimator/amos_helper.py || \
    sed -i 's|from jax.experimental.pjit import PartitionSpec|from jax.sharding import PartitionSpec|g' \
    /usr/lib/python3.11/dist-packages/jestimator/amos_helper.py || true

# Install magenta-realtime without deps
RUN git clone https://github.com/magenta/magenta-realtime.git /magenta-realtime-src && \
    cd /magenta-realtime-src && \
    uv pip install --system --no-deps -e .

# API and audio processing dependencies
RUN uv pip install --system \
    fastapi uvicorn[standard] python-multipart \
    pyloudnorm gradio soxr huggingface_hub

# Ensure compatible protobuf version
RUN uv pip install --system --force-reinstall "protobuf>=5.27.0"

# Create HuggingFace Space user
RUN useradd -m -u 1000 appuser
WORKDIR /home/appuser/app

# Set cache directory
ENV MAGENTA_RT_CACHE_DIR=/home/appuser/.cache/magenta_rt
RUN mkdir -p $MAGENTA_RT_CACHE_DIR && chown -R appuser:appuser /home/appuser/.cache

# Copy application files with proper ownership
COPY --chown=appuser:appuser app.py /home/appuser/app/
COPY --chown=appuser:appuser utils.py /home/appuser/app/
COPY --chown=appuser:appuser jam_worker.py /home/appuser/app/
COPY --chown=appuser:appuser one_shot_generation.py /home/appuser/app/
COPY --chown=appuser:appuser model_management.py /home/appuser/app/
COPY --chown=appuser:appuser documentation.html /home/appuser/app/
COPY --chown=appuser:appuser lil_demo_540p.mp4 /home/appuser/app/
COPY --chown=appuser:appuser magentaRT_rt_tester.html /home/appuser/app/
COPY --chown=appuser:appuser magenta_prompts.js /home/appuser/app/
COPY --chown=appuser:appuser docs/ /home/appuser/app/docs/

USER appuser

EXPOSE 7860
CMD ["bash", "-lc", "python -m uvicorn app:app --host 0.0.0.0 --port ${PORT:-7860}"]