# ====================
# Stage 1: Build nsjail from source
# ====================
FROM debian:bookworm-slim AS nsjail-builder

ENV DEBIAN_FRONTEND=noninteractive
# Build from specific commit that includes pasta/user_net support
ENV NSJAIL_COMMIT=b24be32d38a26656568491c2c5fcffa6e77341d6

RUN apt-get update && apt-get install -y --no-install-recommends \
    git gcc g++ make pkg-config bison flex \
    libprotobuf-dev protobuf-compiler libnl-route-3-dev ca-certificates \
    && rm -rf /var/lib/apt/lists/*

RUN git clone https://github.com/google/nsjail.git /tmp/nsjail && \
    cd /tmp/nsjail && git checkout "${NSJAIL_COMMIT}" && \
    git submodule update --init --recursive && \
    make -j"$(nproc)" && \
    install -m 0755 nsjail /usr/local/bin/nsjail && \
    rm -rf /tmp/nsjail

# ====================
# Stage 2: Create minimal sandbox rootfs
# ====================
FROM node:22.13.1-slim AS node-bin
FROM python:3.12-slim-bookworm AS sandbox-rootfs

ARG TARGETARCH
ARG DUCKDB_VERSION=1.4.3

RUN apt-get update && apt-get install -y --no-install-recommends \
    ca-certificates curl wget jq iputils-ping && rm -rf /var/lib/apt/lists/*

# This rootfs is shared by run_python and agent sandboxes; CLI additions here
# are intentionally available to both. DuckDB is not available from Bookworm
# apt, so install the official release binary, verify it, preinstall extensions,
# and wrap the CLI to load those extensions on each invocation.
RUN arch="${TARGETARCH:-$(dpkg --print-architecture)}" && \
    case "${arch}" in \
        amd64) duckdb_sha256="c479794045d094058d3092e404e696508d6310b5d234a8c1945b745678f09d8d" ;; \
        arm64) duckdb_sha256="c709eb3efc74a609af4b92bc885c509a1bd21ddfa71ea1e717420d4dd9fc121b" ;; \
        *) echo "Unsupported DuckDB CLI architecture: ${arch}" >&2; exit 1 ;; \
    esac && \
    curl -fsSL "https://github.com/duckdb/duckdb/releases/download/v${DUCKDB_VERSION}/duckdb_cli-linux-${arch}.gz" -o /tmp/duckdb.gz && \
    echo "${duckdb_sha256}  /tmp/duckdb.gz" | sha256sum -c - && \
    gunzip /tmp/duckdb.gz && \
    install -m 0755 /tmp/duckdb /usr/local/bin/duckdb.real && \
    rm -f /tmp/duckdb && \
    mkdir -p /usr/local/lib/duckdb/extensions /usr/local/share/duckdb && \
    /usr/local/bin/duckdb.real -c "SET extension_directory = '/usr/local/lib/duckdb/extensions'; INSTALL json; INSTALL postgres; INSTALL httpfs; INSTALL sqlite; INSTALL inet;" && \
    printf '%s\n' \
        "SET extension_directory = '/usr/local/lib/duckdb/extensions';" \
        "LOAD json;" \
        "LOAD postgres;" \
        "LOAD httpfs;" \
        "LOAD sqlite;" \
        "LOAD inet;" \
        > /usr/local/share/duckdb/tracecat-init.sql && \
    printf '%s\n' \
        '#!/bin/sh' \
        'exec /usr/local/bin/duckdb.real -init /usr/local/share/duckdb/tracecat-init.sql "$@"' \
        > /usr/local/bin/duckdb && \
    chmod 0755 /usr/local/bin/duckdb && \
    jq --version && \
    duckdb --version && \
    test "$(duckdb -csv -noheader -c "SELECT count(*) FROM duckdb_extensions() WHERE extension_name IN ('json', 'postgres_scanner', 'httpfs', 'sqlite_scanner', 'inet') AND installed AND loaded;")" = "5"

COPY --from=ghcr.io/astral-sh/uv:0.9.15 /uv /usr/local/bin/uv
COPY --from=ghcr.io/astral-sh/uv:0.9.15 /uvx /usr/local/bin/uvx
COPY --from=node-bin /usr/local/bin/node /usr/local/bin/node
COPY --from=node-bin /usr/local/lib/node_modules /usr/local/lib/node_modules
RUN ln -s ../lib/node_modules/npm/bin/npm-cli.js /usr/local/bin/npm && \
    ln -s ../lib/node_modules/npm/bin/npx-cli.js /usr/local/bin/npx

RUN useradd -m -u 1000 sandbox && \
    mkdir -p /workspace /work /cache /packages /home/sandbox && \
    chown sandbox:sandbox /workspace /work /cache /packages /home/sandbox

ENV HOME=/tmp PYTHONDONTWRITEBYTECODE=1 PYTHONUNBUFFERED=1

# ====================
# Stage 3: Shared base with runtime dependencies
# ====================
FROM ghcr.io/astral-sh/uv:0.9.15-python3.12-bookworm-slim AS base

ENV HOST=0.0.0.0 PORT=8000

# Copy nsjail binary
COPY --from=nsjail-builder /usr/local/bin/nsjail /usr/local/bin/nsjail

# Copy Node.js + npx for in-process MCP command servers (stdio)
COPY --from=node-bin /usr/local/bin/node /usr/local/bin/node
COPY --from=node-bin /usr/local/lib/node_modules /usr/local/lib/node_modules
RUN ln -s ../lib/node_modules/npm/bin/npm-cli.js /usr/local/bin/npm && \
    ln -s ../lib/node_modules/npm/bin/npx-cli.js /usr/local/bin/npx

# Install runtime packages
RUN apt-get update && apt-get install -y --no-install-recommends \
    acl git openssh-client xmlsec1 libmagic1 curl ca-certificates jq \
    libnl-route-3-200 libprotobuf32 libcap2-bin util-linux \
    passt squashfs-tools \
    && apt-get -y upgrade \
    && apt-get clean && rm -rf /var/lib/apt/lists/*

COPY --from=sandbox-rootfs /usr/local/bin/duckdb /usr/local/bin/duckdb
COPY --from=sandbox-rootfs /usr/local/bin/duckdb.real /usr/local/bin/duckdb.real
COPY --from=sandbox-rootfs /usr/local/lib/duckdb /usr/local/lib/duckdb
COPY --from=sandbox-rootfs /usr/local/share/duckdb /usr/local/share/duckdb

RUN jq --version && duckdb --version

# Allow the non-root executor process to invoke mount/umount when the container
# runtime grants the needed bounding capabilities. Without these file caps,
# privileged Docker containers still run apiuser with no effective capabilities.
RUN chmod u-s /usr/bin/mount /usr/bin/umount && \
    setcap cap_sys_admin,cap_dac_override+ep /usr/bin/mount && \
    setcap cap_sys_admin,cap_dac_override+ep /usr/bin/umount

# Copy sandbox rootfs
COPY --from=sandbox-rootfs /usr /var/lib/tracecat/sandbox-rootfs/usr
COPY --from=sandbox-rootfs /lib /var/lib/tracecat/sandbox-rootfs/lib
COPY --from=sandbox-rootfs /bin /var/lib/tracecat/sandbox-rootfs/bin
COPY --from=sandbox-rootfs /sbin /var/lib/tracecat/sandbox-rootfs/sbin
COPY --from=sandbox-rootfs /etc/passwd /var/lib/tracecat/sandbox-rootfs/etc/passwd
COPY --from=sandbox-rootfs /etc/group /var/lib/tracecat/sandbox-rootfs/etc/group
COPY --from=sandbox-rootfs /etc/ssl /var/lib/tracecat/sandbox-rootfs/etc/ssl
COPY --from=sandbox-rootfs /etc/ca-certificates /var/lib/tracecat/sandbox-rootfs/etc/ca-certificates
RUN install -m 0644 /dev/null /var/lib/tracecat/sandbox-rootfs/etc/resolv.conf && \
    install -m 0644 /dev/null /var/lib/tracecat/sandbox-rootfs/etc/hosts && \
    install -m 0644 /dev/null /var/lib/tracecat/sandbox-rootfs/etc/nsswitch.conf

# Handle lib64 for amd64
RUN if [ -d /var/lib/tracecat/sandbox-rootfs/lib64 ] || [ "$(uname -m)" = "x86_64" ]; then \
        mkdir -p /var/lib/tracecat/sandbox-rootfs/lib64 && \
        cp -a /lib64/. /var/lib/tracecat/sandbox-rootfs/lib64/ 2>/dev/null || true; \
    fi

# Create sandbox directories
RUN mkdir -p /var/lib/tracecat/sandbox-rootfs/tmp \
    /var/lib/tracecat/sandbox-rootfs/proc \
    /var/lib/tracecat/sandbox-rootfs/dev \
    /var/lib/tracecat/sandbox-rootfs/work \
    /var/lib/tracecat/sandbox-rootfs/cache \
    /var/lib/tracecat/sandbox-rootfs/packages \
    /var/lib/tracecat/sandbox-rootfs/home/sandbox \
    /var/lib/tracecat/sandbox-cache/packages \
    /var/lib/tracecat/sandbox-cache/uv-cache && \
    chmod -R 755 /var/lib/tracecat/sandbox-rootfs && \
    chown -R 1000:1000 /var/lib/tracecat/sandbox-rootfs/work \
        /var/lib/tracecat/sandbox-rootfs/cache \
        /var/lib/tracecat/sandbox-rootfs/packages \
        /var/lib/tracecat/sandbox-rootfs/home/sandbox && \
    chmod 1777 /var/lib/tracecat/sandbox-rootfs/tmp

# Create apiuser for non-root runtime (required for pasta userspace networking)
RUN groupadd -g 1001 apiuser && useradd -m -u 1001 -g apiuser apiuser && \
    mkdir -p /home/apiuser/.cache/uv /home/apiuser/.cache/s3 /home/apiuser/.cache/tmp /home/apiuser/.local/bin && \
    chown -R apiuser:apiuser /home/apiuser

# Create MCP socket directory for apiuser
RUN mkdir -p /var/run/tracecat && chown 1001:1001 /var/run/tracecat

WORKDIR /app

# ====================
# Stage 4: Development target
# ====================
FROM base AS development

ENV TMPDIR="/home/apiuser/.cache/tmp" TEMP="/home/apiuser/.cache/tmp" TMP="/home/apiuser/.cache/tmp"

# Set sandbox cache permissions for apiuser
RUN chown -R 1001:1001 /var/lib/tracecat/sandbox-cache && \
    chmod -R 755 /var/lib/tracecat/sandbox-cache

# Prime uv cache (as root, before switching user)
RUN --mount=type=cache,target=/root/.cache/uv \
    --mount=type=bind,source=uv.lock,target=uv.lock \
    --mount=type=bind,source=pyproject.toml,target=pyproject.toml \
    --mount=type=bind,source=packages,target=packages \
    uv sync --locked --no-install-project --no-dev --no-editable

COPY --chown=apiuser:apiuser . /app/

RUN --mount=type=cache,target=/root/.cache/uv uv sync --frozen --no-dev

# Fix ownership of /app (uv sync creates .venv as root)
RUN chown -R apiuser:apiuser /app

ENV PATH="/app/.venv/bin:$PATH"
ENV PYTHONPATH="/home/apiuser/.local"

RUN mkdir -p /home/apiuser/.local/bin && ln -s $(which uv) /home/apiuser/.local/bin/uv

COPY docker/scripts/entrypoint.sh /app/entrypoint.sh
RUN chmod +x /app/entrypoint.sh

# Switch to non-root user (matches production, required for pasta userspace networking)
USER apiuser

ENTRYPOINT ["/app/entrypoint.sh"]
EXPOSE $PORT
CMD ["sh", "-c", "python3 -m uvicorn tracecat.api.app:app --host $HOST --port $PORT --reload"]

# ====================
# Stage 5: Test target (development + pytest)
# ====================
FROM development AS test

# Install test dependencies
RUN --mount=type=cache,target=/home/apiuser/.cache/uv,uid=1001,gid=1001 uv sync --frozen --group dev

CMD ["python", "-m", "pytest"]

# ====================
# Stage 6: Production target
# ====================
FROM base AS production

ENV UV_COMPILE_BYTECODE=1 UV_LINK_MODE=copy

# Set sandbox cache permissions for apiuser
RUN chown -R 1001:1001 /var/lib/tracecat/sandbox-cache && \
    chmod -R 755 /var/lib/tracecat/sandbox-cache

ENV PYTHONUSERBASE="/home/apiuser/.local"
ENV UV_CACHE_DIR="/home/apiuser/.cache/uv"
ENV PYTHONPATH="/home/apiuser/.local"
ENV PATH="/home/apiuser/.local/bin:/usr/local/bin:/usr/bin:/bin"
ENV TMPDIR="/home/apiuser/.cache/tmp" TEMP="/home/apiuser/.cache/tmp" TMP="/home/apiuser/.cache/tmp"

RUN mkdir -p /app/.scripts && chown -R apiuser:apiuser /app

# Switch to non-root user
USER apiuser

# Install dependencies as apiuser
RUN --mount=type=cache,target=/home/apiuser/.cache/uv,uid=1001,gid=1001 \
    --mount=type=bind,source=uv.lock,target=uv.lock \
    --mount=type=bind,source=pyproject.toml,target=pyproject.toml \
    --mount=type=bind,source=packages,target=packages \
    uv sync --locked --no-install-project --no-dev --no-editable

COPY --chown=apiuser:apiuser ./tracecat /app/tracecat
COPY --chown=apiuser:apiuser ./packages /app/packages
COPY --chown=apiuser:apiuser ./pyproject.toml ./uv.lock ./.python-version ./README.md ./LICENSE ./alembic.ini /app/
COPY --chown=apiuser:apiuser ./alembic /app/alembic
COPY --chown=apiuser:apiuser docker/scripts/entrypoint.sh /app/entrypoint.sh
RUN chmod +x /app/entrypoint.sh

RUN --mount=type=cache,target=/home/apiuser/.cache/uv,uid=1001,gid=1001 \
    uv sync --locked --no-dev --no-editable

ENV PATH="/app/.venv/bin:/home/apiuser/.local/bin:/usr/local/bin:/usr/bin:/bin"

RUN ln -sf $(which uv) /home/apiuser/.local/bin/uv

# Verification
RUN nsjail --help > /dev/null 2>&1 && echo "nsjail available"

EXPOSE $PORT
ENTRYPOINT ["/app/entrypoint.sh"]
CMD ["sh", "-c", "python3 -m uvicorn tracecat.api.app:app --host $HOST --port $PORT"]
