# Dockerfile for HuggingFace Space deployment (x86_64)
# Based on working Modal Dockerfile.x86
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="cuda"

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
RUN uv pip install --system tf-nightly

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

# -------------------------------------------------------------------
# HARD STABILIZER: prevent resampy->numba from selecting py<3.10-only wheels
# Also keep numpy < 2 to avoid later numba incompat headaches on py311.
# -------------------------------------------------------------------
RUN printf "numpy<2\nnumba>=0.57.0\nllvmlite>=0.40.0\n" > /tmp/constraints.txt && \
    echo "---- constraints.txt ----" && \
    cat /tmp/constraints.txt && \
    echo "------------------------"

# (Optional but often helpful) install pinned trio first so they land as wheels cleanly,
# then everything else resolves around them.
RUN uv pip install --system -c /tmp/constraints.txt "numpy<2" "numba>=0.57.0" "llvmlite>=0.40.0"

# Install base dependencies (constraints applied)
RUN uv pip install --system -c /tmp/constraints.txt \
    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 . && \
    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 (constraints applied)
RUN uv pip install --system -c /tmp/constraints.txt \
    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 (constraints applied)
RUN uv pip install --system -c /tmp/constraints.txt \
    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"

RUN uv pip install --system fastrtc

# Set working directory and create cache
WORKDIR /app
ENV MAGENTA_RT_CACHE_DIR=/root/.cache/magenta_rt
RUN mkdir -p $MAGENTA_RT_CACHE_DIR

# Copy application files for HuggingFace Space
COPY app.py /app/
COPY utils.py /app/
COPY jam_worker.py /app/
COPY one_shot_generation.py /app/
COPY model_management.py /app/
COPY documentation.html /app/
COPY lil_demo_540p.mp4 /app/
COPY magentaRT_rt_tester.html /app/
COPY magenta_prompts.js /app/
COPY docs/ /app/docs/
COPY fastrtc_magenta.py /app/

EXPOSE 7860
CMD ["python", "-m", "uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]