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

ROOT=$(CDPATH= cd -- "$(dirname -- "$0")/.." && pwd)
CLI=${CLI:-$ROOT/.venv/bin/unlimited-ocr-rdna4}
PYTHON=${PYTHON:-$ROOT/.venv/bin/python}
DEVICE=${UNLIMITED_OCR_DEVICE:-0}
MODEL_DIR=${UNLIMITED_OCR_MODEL_DIR:-}
AMD_SMI_GPU=${AMD_SMI_GPU:-}
NETWORK_ISOLATED=${VALIDATION_NETWORK_ISOLATED:-0}
EXPECTED_IMAGE_SHA256=${EXPECTED_IMAGE_SHA256:-13df40051aceb3bd14f75622697dda1f2a96515bb95c0483e0920cd4d41cbc01}
EXPECTED_PDF_SHA256=${EXPECTED_PDF_SHA256:-fb17a6396c3e73a025e388a09f73357960ca0bf95124e38179e06a2e64442438}
EVIDENCE_OUTPUT=${VALIDATION_EVIDENCE_OUTPUT:-}
CACHE_BASE=${XDG_CACHE_HOME:-${HOME}/.cache}/unlimited-ocr-rdna4
FIXTURE_IMAGE="$ROOT/tests/fixtures/rdna4-smoke.png"
FIXTURE_PDF="$ROOT/tests/fixtures/rdna4-smoke.pdf"
mkdir -p "$CACHE_BASE"
WORK_DIR=$(mktemp -d "$CACHE_BASE/validation.XXXXXX")
trap 'rm -rf -- "$WORK_DIR"' EXIT

printf '%s  %s\n' \
  'f5099e17be868abfb4213dbdab220deac82a2db93ba87ab22f03219178246972' "$FIXTURE_IMAGE" \
  'cdd2b484d0ac90bd98b489dd97565a65eb17246f713359524cddd41d78cc10cb' "$FIXTURE_PDF" |
  sha256sum --check --status

MODEL_ARGS=()
if [[ -n $MODEL_DIR ]]; then
  MODEL_ARGS=(--model-dir "$MODEL_DIR")
fi

run_isolated() {
  sudo -n systemd-run \
    --quiet --wait --pipe --collect \
    -p PrivateNetwork=yes \
    -p "User=$(id -un)" \
    -p "WorkingDirectory=$ROOT" \
    -p "Environment=HOME=$HOME" \
    -- "$@"
}

run_ocr() {
  if [[ $NETWORK_ISOLATED == 1 ]]; then
    run_isolated "$@"
  else
    "$@"
  fi
}

if [[ $NETWORK_ISOLATED == 1 ]]; then
  if run_isolated /usr/bin/python3 -c \
    'import socket; socket.create_connection(("1.1.1.1", 80), timeout=0.25)' \
    >"$WORK_DIR/network-probe.out" 2>"$WORK_DIR/network-probe.err"; then
    printf 'error: PrivateNetwork validation probe unexpectedly reached the network\n' >&2
    exit 5
  fi
fi

"$CLI" doctor --device "$DEVICE" --require-model --json "${MODEL_ARGS[@]}" >"$WORK_DIR/doctor.json"

run_ocr "$CLI" run \
  --input "$FIXTURE_IMAGE" \
  --output "$WORK_DIR/smoke-image-1.md" \
  --device "$DEVICE" \
  --json \
  "${MODEL_ARGS[@]}" >"$WORK_DIR/image-1.json"

run_ocr "$CLI" run \
  --input "$FIXTURE_IMAGE" \
  --output "$WORK_DIR/smoke-image-2.md" \
  --device "$DEVICE" \
  --json \
  "${MODEL_ARGS[@]}" >"$WORK_DIR/image-2.json"

run_ocr "$CLI" run \
  --input "$FIXTURE_PDF" \
  --output "$WORK_DIR/smoke-pdf.md" \
  --device "$DEVICE" \
  --dpi 150 \
  --max-pages 1 \
  --json \
  "${MODEL_ARGS[@]}" >"$WORK_DIR/pdf.json"

if [[ -n $AMD_SMI_GPU ]]; then
  amd-smi process -g "$AMD_SMI_GPU" --json >"$WORK_DIR/gpu-processes.json"
fi

CHECK_ARGS=(
  --doctor "$WORK_DIR/doctor.json"
  --image-run "$WORK_DIR/image-1.json"
  --image-run "$WORK_DIR/image-2.json"
  --pdf-run "$WORK_DIR/pdf.json"
  --image-output "$WORK_DIR/smoke-image-1.md"
  --image-output "$WORK_DIR/smoke-image-2.md"
  --pdf-output "$WORK_DIR/smoke-pdf.md"
)
CHECK_ARGS+=(--expected-image-sha256 "$EXPECTED_IMAGE_SHA256")
CHECK_ARGS+=(--expected-pdf-sha256 "$EXPECTED_PDF_SHA256")
if [[ $NETWORK_ISOLATED == 1 ]]; then
  CHECK_ARGS+=(--network-isolated)
fi
if [[ -n $AMD_SMI_GPU ]]; then
  CHECK_ARGS+=(--gpu-processes "$WORK_DIR/gpu-processes.json")
fi
if [[ -n $EVIDENCE_OUTPUT ]]; then
  CHECK_ARGS+=(--evidence-output "$EVIDENCE_OUTPUT")
fi

"$PYTHON" "$ROOT/scripts/check-validation.py" "${CHECK_ARGS[@]}"