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

COMPOSE="${COMPOSE:-docker compose -f docker-compose.test.yml --profile jetson}"
REGISTRY_HOST="${REGISTRY_HOST:?REGISTRY_HOST is required}"
ENGINE_REPOSITORY="${JETSON_ENGINE_REPOSITORY:-$REGISTRY_HOST/azaion/detections-jetson-engine}"
BRANCH="${CI_COMMIT_BRANCH:-local}"
default_tag="$(printf '%s' "$BRANCH" | sed 's#^refs/heads/##; s#^refs/tags/##; s#[^A-Za-z0-9_.-]#-#g; s#^-*##; s#-*$##')"
default_tag="${default_tag:-local}"
ENGINE_TAG="${JETSON_ENGINE_TAG:-$default_tag}"
OUT_DIR="${JETSON_ENGINE_OUT_DIR:-results/jetson-engine}"

mkdir -p "$OUT_DIR/models"

loader_id="$($COMPOSE ps -q mock-loader)"
if [[ -z "$loader_id" ]]; then
  echo "ERROR: mock-loader container is not running"
  exit 1
fi

docker cp "$loader_id:/models/models/." "$OUT_DIR/models/"
find "$OUT_DIR/models" -maxdepth 1 -type f ! -name 'azaion*.engine' -delete

engine_count="$(find "$OUT_DIR/models" -maxdepth 1 -type f -name 'azaion*.engine' | wc -l | tr -d ' ')"
if [[ "$engine_count" == "0" ]]; then
  echo "ERROR: no converted TensorRT engine found in mock-loader /models/models"
  find "$OUT_DIR/models" -maxdepth 2 -type f -print
  exit 1
fi

if [[ "${REQUIRE_INT8_ENGINE:-0}" == "1" ]]; then
  int8_engine_count="$(find "$OUT_DIR/models" -maxdepth 1 -type f -name 'azaion*.int8.engine' | wc -l | tr -d ' ')"
  if [[ "$int8_engine_count" == "0" ]]; then
    echo "ERROR: INT8 engine is required, but no azaion*.int8.engine file was produced"
    find "$OUT_DIR/models" -maxdepth 1 -type f -name 'azaion*.engine' -print -exec ls -lh {} \;
    exit 1
  fi
fi

echo "--- Converted TensorRT engine files:"
find "$OUT_DIR/models" -maxdepth 1 -type f -name 'azaion*.engine' -print -exec ls -lh {} \;

image="$ENGINE_REPOSITORY:$ENGINE_TAG"
echo "--- Building Jetson engine artifact image: $image"
docker build -f engine-artifact.Dockerfile -t "$image" "$OUT_DIR/models"
docker push "$image"

if [[ -n "${CI_COMMIT_SHA:-}" ]]; then
  sha_tag="$(printf '%s' "$CI_COMMIT_SHA" | cut -c1-12)"
  docker tag "$image" "$ENGINE_REPOSITORY:$sha_tag"
  docker push "$ENGINE_REPOSITORY:$sha_tag"
fi

echo "--- Published Jetson engine artifact image: $image"