#!/usr/bin/env bash
set -euo pipefail

MODEL_REPOSITORY="fastino/gliner2-privacy-filter-PII-multi"
MODEL_REVISION="c153999da5f4c509df4322b0c6a1baf3d2c284d7"
TARGET="${1:-}"
ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd -P)"

if [[ -z "$TARGET" || "$TARGET" != /* ]]; then
  echo "Usage: $0 /absolute/model/directory" >&2
  exit 2
fi

if [[ -d "$TARGET" ]] && find "$TARGET" -mindepth 1 -print -quit | grep -q .; then
  echo "Target directory must be empty: $TARGET" >&2
  exit 2
fi
mkdir -p "$TARGET"
cd "$ROOT"

docker build \
  --build-arg INSTALL_SENSITIVITY_NER=true \
  --file docker/core.Dockerfile \
  --tag thothii-core:sensitivity-ner \
  .

docker run --rm \
  --user "$(id -u):$(id -g)" \
  --env HOME=/tmp \
  --env TMPDIR=/tmp \
  --env XDG_CACHE_HOME=/tmp/.cache \
  --env HF_HOME=/tmp/huggingface \
  --env HF_HUB_DISABLE_TELEMETRY=1 \
  --env HF_HUB_DISABLE_XET=1 \
  --volume "$TARGET:/model" \
  --entrypoint /opt/sensitivity-ner/bin/hf \
  thothii-core:sensitivity-ner \
  download "$MODEL_REPOSITORY" \
  --revision "$MODEL_REVISION" \
  --local-dir /model

printf '%s\n' "$MODEL_REVISION" > "$TARGET/THOTHII_MODEL_REVISION"

docker run --rm \
  --user "$(id -u):$(id -g)" \
  --volume "$TARGET:/model" \
  --entrypoint /bin/sh \
  thothii-core:sensitivity-ner \
  -c 'cd /model && sha256sum -c /app/backend/python/sensitivity-ner-model-sha256.txt && cp /app/backend/python/sensitivity-ner-model-sha256.txt MODEL_SHA256SUMS'

docker run --rm \
  --network none \
  --entrypoint /opt/sensitivity-ner/bin/pip \
  thothii-core:sensitivity-ner \
  check

GPU_AUDIT="$(docker run --rm \
  --network none \
  --entrypoint /opt/sensitivity-ner/bin/python \
  thothii-core:sensitivity-ner \
  -c 'import torch; print(f"{torch.cuda.is_available()}|{torch.version.cuda}|{torch.cuda.device_count()}")')"
if [[ "$GPU_AUDIT" != "False|None|0" ]]; then
  echo "The optional runtime is not CPU-only: $GPU_AUDIT" >&2
  exit 1
fi

NETWORK_GUARD_AUDIT="$(docker run --rm \
  --entrypoint /opt/sensitivity-ner/bin/python \
  thothii-core:sensitivity-ner \
  -c '
import importlib.util
import socket
spec = importlib.util.spec_from_file_location("worker", "/app/backend/python/sensitivity_ner_worker.py")
worker = importlib.util.module_from_spec(spec)
assert spec.loader is not None
spec.loader.exec_module(worker)
raw_socket = socket.socket
worker._disable_network()
try:
    raw_socket(socket.AF_INET, socket.SOCK_STREAM)
except PermissionError as error:
    print(error.errno)
else:
    raise SystemExit("network syscall filter is inactive")
')"
if [[ "$NETWORK_GUARD_AUDIT" != "1" ]]; then
  echo "The optional runtime did not activate its network syscall filter" >&2
  exit 1
fi

SMOKE_OUTPUT="$(mktemp)"
SMOKE_ERROR="$(mktemp)"
trap 'rm -f -- "$SMOKE_OUTPUT" "$SMOKE_ERROR"' EXIT
printf '%s\n' '{"id":"smoke","candidates":[{"columnId":"33333333-3333-4333-8333-333333333333","text":"La paziente si chiama Maria Rossi."}]}' \
  | docker run --rm \
      --network none \
      --read-only \
      --tmpfs /tmp:rw,noexec,nosuid,size=512m \
      --volume "$TARGET:/model:ro" \
      --entrypoint /opt/sensitivity-ner/bin/python \
      -i thothii-core:sensitivity-ner \
      /app/backend/python/sensitivity_ner_worker.py --model /model --threads 2 \
      >"$SMOKE_OUTPUT" 2>"$SMOKE_ERROR"
if ! grep -Fxq '{"ready":true}' "$SMOKE_OUTPUT" \
  || ! grep -Eq '"ok":true.*"columnId":"33333333-3333-4333-8333-333333333333".*"label":"full_name"' "$SMOKE_OUTPUT"; then
  echo "The offline Italian CPU smoke test failed" >&2
  sed -n '1,20p' "$SMOKE_ERROR" >&2
  exit 1
fi

echo "Pinned model downloaded to $TARGET"