#!/bin/bash

set -euo pipefail

if [ ! -d "$1" ]; then
  echo "Usage: $0 [model directory]" >&2
  exit 1
fi

if [ ! -f "$1/config.json" ] || [ ! -f "$1/model.pt" ]; then
  echo "Model or config not found in $1" >&2
  exit 1
fi

dname=$(basename "$1")

aws s3 sync "$1" "s3://wavtool-composer-models/$dname" --exclude "*" --include "model.pt" --include "config.json"

aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 035363416972.dkr.ecr.us-east-1.amazonaws.com
docker buildx build . -f Dockerfile.tensorrt_builder -t 035363416972.dkr.ecr.us-east-1.amazonaws.com/composer-bundler --push

jobname="tensorrt-builder-$(openssl rand -hex 4)"
kubectl run "$jobname" \
  -i --rm \
  --restart=Never \
  --image=035363416972.dkr.ecr.us-east-1.amazonaws.com/composer-bundler \
  --pod-running-timeout=10m \
  --override-type strategic \
  --overrides="$(cat <<EOF
{
  "apiVersion": "v1",
  "spec": {
    "serviceAccountName": "composer-service-account",
    "nodeSelector": {
      "karpenter.k8s.aws/instance-family": "g4dn",
      "karpenter.sh/capacity-type": "on-demand"
    },
    "containers": [
      {
        "name": "$jobname",
        "resources": {
          "limits": {
            "nvidia.com/gpu": "1"
          }
        }
      }
    ]
  }
}
EOF
)" \
  --env="MODEL_NAME=$dname" \
  --command -- bash -c "$(cat <<'EOF'
set -euo pipefail

aws s3 cp "s3://wavtool-composer-models/$MODEL_NAME/model.pt" ./model.pt
aws s3 cp "s3://wavtool-composer-models/$MODEL_NAME/config.json" ./config.json

. /opt/venv/bin/activate

python onnx_export.py \
  --config-path=config.json \
  --model-path=model.pt \
  export model.onnx model_one_step.onnx

python trt_builder.py \
  --config_path=config.json \
  build-trt model.onnx model.trt

python trt_builder.py \
  --config_path=config.json \
  build-trt-one-step model_one_step.onnx model_one_step.trt

aws s3 cp model.trt "s3://wavtool-composer-models/$MODEL_NAME/model.trt"
aws s3 cp model_one_step.trt "s3://wavtool-composer-models/$MODEL_NAME/model_one_step.trt"

EOF
)"

aws s3 cp "s3://wavtool-composer-models/$dname/model.trt" "$1/model.trt"
aws s3 cp "s3://wavtool-composer-models/$dname/model_one_step.trt" "$1/model_one_step.trt"
aws s3 rm "s3://wavtool-composer-models/$dname/model.pt"
aws s3 rm "s3://wavtool-composer-models/$dname/config.json"