# =============================================================================
# StockGNN R9700 — ROCm Ubuntu Docker Image
# =============================================================================
# Base image: Ubuntu 22.04 with ROCm 5.6 pre-installed
# Provides AMD GPU support for PyTorch/ PyG on Radeon R9700 AI Pro
# =============================================================================

FROM rocm/dev-ubuntu-22.04:5.6

LABEL maintainer="StockGNN Team"
LABEL description="AMD-optimized GNN trading system (ROCm 5.6 / Ubuntu 22.04)"

# ---------------------------------------------------------------------------
# System dependencies
# ---------------------------------------------------------------------------
ENV DEBIAN_FRONTEND=noninteractive
RUN apt-get update && apt-get install -y --no-install-recommends \
    python3.10 \
    python3.10-dev \
    python3-pip \
    git \
    wget \
    curl \
    ca-certificates \
    libglib2.0-0 \
    libsm6 \
    libxext6 \
    libxrender-dev \
    libgomp1 \
    libnuma1 \
    # build tools required for compiling PyG extensions from source
    build-essential \
    cmake \
    ninja-build \
    && rm -rf /var/lib/apt/lists/*

# Make python3.10 the default python3
RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.10 1 \
    && update-alternatives --install /usr/bin/python python /usr/bin/python3.10 1

# Upgrade pip
RUN python -m pip install --no-cache-dir --upgrade pip setuptools wheel

# ---------------------------------------------------------------------------
# Install PyTorch with ROCm 5.6 support (separate from other deps)
# ---------------------------------------------------------------------------
RUN pip install --no-cache-dir \
    torch==2.1.0+rocm5.6 \
    torchvision==0.16.0+rocm5.6 \
    torchaudio==2.1.0+rocm5.6 \
    --index-url https://download.pytorch.org/whl/rocm5.6

# ---------------------------------------------------------------------------
# Install PyTorch Geometric and extensions
# ---------------------------------------------------------------------------
# ROCm 5.6 does not ship prebuilt PyG extension wheels, so we compile them
# from source.  torch-geometric itself is pure-Python and installs from PyPI.
# The extensions (scatter, sparse, cluster, spline-conv) are built with
# torch.utils.cpp_extension which automatically picks up the ROCm toolchain
# shipped in the base image.
# ---------------------------------------------------------------------------
ENV ROCM_HOME=/opt/rocm
ENV HIP_HOME=/opt/rocm

RUN pip install --no-cache-dir torch-geometric==2.4.0 pybind11

# Build PyG extensions with no build isolation so they can see torch
RUN pip install --no-cache-dir --no-build-isolation \
    torch-scatter==2.1.2 \
    torch-sparse==0.6.18 \
    torch-cluster==1.6.2 \
    torch-spline-conv==1.2.2

# ---------------------------------------------------------------------------
# Install remaining Python dependencies from PyPI
# ---------------------------------------------------------------------------
COPY requirements.txt /tmp/requirements.txt

# Strip out the torch / pyg lines and ROCm-specific index directives so we
# install the rest of the packages from standard PyPI.
RUN grep -vE '^(torch|torchvision|torchaudio|torch-geometric|torch-scatter|torch-sparse|torch-cluster|torch-spline-conv|rocblas|hipblaslt|miopen-hip|rccl|--index-url|--find-links)' /tmp/requirements.txt \
    > /tmp/requirements-clean.txt || true

RUN pip install --no-cache-dir -r /tmp/requirements-clean.txt

# ---------------------------------------------------------------------------
# ROCm runtime tuning for AMD Radeon R9700 AI Pro (gfx1030 / gfx1100)
# ---------------------------------------------------------------------------
ENV HSA_OVERRIDE_GFX_VERSION=10.3.0
ENV PYTORCH_HIP_ALLOC_CONF=expandable_segments:True
ENV ROCM_PATH=/opt/rocm
ENV PATH="${ROCM_PATH}/bin:${PATH}"
ENV LD_LIBRARY_PATH="${ROCM_PATH}/lib:${LD_LIBRARY_PATH}"

# ---------------------------------------------------------------------------
# Application setup
# ---------------------------------------------------------------------------
WORKDIR /app

# Create directories that the app expects
RUN mkdir -p /app/data/raw /app/data/processed /app/data/external /app/models /app/logs

# Copy source code
COPY . /app/

# Ensure Python can find local src/ modules
ENV PYTHONPATH="/app:${PYTHONPATH}"

# Create non-root user for runtime security
RUN groupadd -r trader && useradd -r -g trader -d /app trader \
    && chown -R trader:trader /app

USER trader

# Default command (can be overridden per-service in docker-compose.yml)
CMD ["python", "-m", "uvicorn", "src.web.app:app", "--host", "0.0.0.0", "--port", "8000"]