# ===========================================================================
# Trainything — API server + GPU runtime
# ===========================================================================
# Single container: serves the REST/WebSocket API and runs ML pipelines.
# The frontend is hosted separately (e.g. GitHub Pages) and connects directly.
#
# Works with Docker and Podman (standard OCI image).
#
# Quick start:
#   docker run --gpus all -p 8000:8000 -e TRAINYTHING_TOKEN=mysecret trainything
#
# Environment variables:
#   TRAINYTHING_TOKEN        — Bearer token for API auth (empty = no auth)
#   TRAINYTHING_IDLE_TIMEOUT — Auto-shutdown after N seconds of inactivity (0 = disabled)
#   TRAINYTHING_DATA_DIR     — Data directory (default: /data)
#
# ===========================================================================

FROM python:3.11-slim

ARG TORCH_INDEX=https://download.pytorch.org/whl/cu124
ARG VERSION=dev
LABEL org.opencontainers.image.version="${VERSION}"
LABEL org.opencontainers.image.source="https://github.com/kvark/trainything"

RUN apt-get update && apt-get install -y --no-install-recommends \
        git curl build-essential linux-libc-dev \
        libegl1 libgl1 libosmesa6-dev \
    && rm -rf /var/lib/apt/lists/*

WORKDIR /app

# PyTorch first (large, changes rarely — good cache layer)
RUN pip install --no-cache-dir \
    torch torchvision --index-url ${TORCH_INDEX}

# ---- Batteries included: simulation & robotics ---------------------------
# These are pre-installed so users can build simulation pipelines out of the
# box (MuJoCo + LeRobot environments, SmolVLA, diffusion models, etc.)
RUN pip install --no-cache-dir \
    mujoco>=3.1 \
    gymnasium>=1.1 \
    gymnasium-robotics>=1.2 \
    stable-retro>=0.9 \
    lerobot>=0.1 \
    diffusers>=0.27 \
    accelerate>=0.29
# ---------------------------------------------------------------------------

# swig is needed to build Box2D wheels
RUN apt-get update && apt-get install -y --no-install-recommends swig \
    && rm -rf /var/lib/apt/lists/*

# Trainything package (with all optional extras)
COPY pyproject.toml README.md ./
COPY server/ ./server/
RUN pip install --no-cache-dir ".[all]"

# Example pipelines and tests
COPY examples/ ./examples/
COPY tests/ ./tests/

# Data directory
ENV TRAINYTHING_DATA_DIR=/data
RUN mkdir -p /data

EXPOSE 8000

ENTRYPOINT ["trainything"]
CMD ["serve", "--host", "0.0.0.0", "--port", "8000", "--tls"]
