# syntax=docker/dockerfile:1
ARG ROCM_VERSION=6.4.1
ARG ROCM_PYTORCH_RELEASE=2.6.0
ARG SPACY_MODEL_URL=https://github.com/explosion/spacy-models/releases/download/en_core_web_sm-3.8.0/en_core_web_sm-3.8.0-py3-none-any.whl

# =============================================================================
# Base Stage - ROCm PyTorch image
# =============================================================================
FROM rocm/pytorch:rocm${ROCM_VERSION}_ubuntu24.04_py3.12_pytorch_release_${ROCM_PYTORCH_RELEASE} AS base

LABEL maintainer="NoteFlow Team"
LABEL description="NoteFlow with ROCm GPU support for AMD GPUs"

ENV PYTHONDONTWRITEBYTECODE=1 \
    PYTHONUNBUFFERED=1 \
    UV_COMPILE_BYTECODE=1 \
    UV_LINK_MODE=copy \
    PATH=/usr/local/bin:$PATH

COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/

RUN apt-get update && apt-get install -y --no-install-recommends \
    build-essential \
    pkg-config \
    python3-venv \
    portaudio19-dev \
    libsndfile1 \
    ffmpeg \
    git \
    && rm -rf /var/lib/apt/lists/*

WORKDIR /workspace

COPY pyproject.toml uv.lock* ./
COPY README.md ./
COPY src ./src/

ENV VIRTUAL_ENV=/opt/venv
RUN uv venv --system-site-packages ${VIRTUAL_ENV}
ENV PATH="${VIRTUAL_ENV}/bin:$PATH"

# =============================================================================
# Production Server Stage (minimal deps)
# =============================================================================
FROM base AS server
ARG SPACY_MODEL_URL

RUN uv pip install --python ${VIRTUAL_ENV}/bin/python -e ".[rocm]"
RUN uv pip install --python ${VIRTUAL_ENV}/bin/python ${SPACY_MODEL_URL}

COPY . .

ENV ROCM_PATH=/opt/rocm \
    HIP_VISIBLE_DEVICES=0 \
    HSA_OVERRIDE_GFX_VERSION="" \
    NOTEFLOW_ASR_DEVICE=rocm \
    NOTEFLOW_FEATURE_ROCM_ENABLED=true

RUN useradd --create-home --shell /bin/bash noteflow
USER noteflow

EXPOSE 50051

HEALTHCHECK --interval=30s --timeout=10s --start-period=5s --retries=3 \
    CMD python -c "import grpc; channel = grpc.insecure_channel('localhost:50051'); grpc.channel_ready_future(channel).result(timeout=5)" || exit 1

CMD ["python", "-m", "noteflow.grpc.server"]

# =============================================================================
# Development Server Stage (all extras + hot reload)
# =============================================================================
FROM base AS dev
ARG SPACY_MODEL_URL

RUN uv pip install --python ${VIRTUAL_ENV}/bin/python -e ".[rocm,optional]"
RUN uv pip install --python ${VIRTUAL_ENV}/bin/python watchfiles
RUN uv pip install --python ${VIRTUAL_ENV}/bin/python ${SPACY_MODEL_URL}

COPY . .

ENV ROCM_PATH=/opt/rocm \
    HIP_VISIBLE_DEVICES=0 \
    HSA_OVERRIDE_GFX_VERSION="" \
    NOTEFLOW_ASR_DEVICE=rocm \
    NOTEFLOW_FEATURE_ROCM_ENABLED=true

EXPOSE 50051

CMD ["python", "scripts/dev_watch_server.py"]