#!/usr/bin/env bash
# Clear LiteLLM Postgres spend/logs on IFT (keep schema + keys if present)
set -euo pipefail
ssh -o BatchMode=yes eventhub-ift bash -s <<'REMOTE'
set -euo pipefail
cd /opt/ai-router-stack
set -a; source .env; set +a
PG=$(docker ps -q -f name=ai-router_postgres | head -1)
if [[ -z "$PG" ]]; then
  echo "ERROR: postgres container not running" >&2
  exit 1
fi
USER="${POSTGRES_USER:-litellm}"
DB="${POSTGRES_DB:-litellm}"
echo "== Tables before =="
docker exec -e PGPASSWORD="$POSTGRES_PASSWORD" "$PG" \
  psql -U "$USER" -d "$DB" -c "\dt"

# LiteLLM Prisma tables — truncate data, keep schema
# Prefer spend/usage; also clear invite/audit-ish if present
SQL=$(cat <<'EOS'
DO $$
DECLARE
  r RECORD;
  n bigint;
BEGIN
  FOR r IN
    SELECT tablename
    FROM pg_tables
    WHERE schemaname = 'public'
      AND tablename NOT IN ('_prisma_migrations')
  LOOP
    EXECUTE format('SELECT count(*) FROM %I', r.tablename) INTO n;
    RAISE NOTICE 'truncate % (% rows)', r.tablename, n;
    EXECUTE format('TRUNCATE TABLE %I RESTART IDENTITY CASCADE', r.tablename);
  END LOOP;
END $$;
EOS
)

echo "== Truncate all public tables (except _prisma_migrations) =="
docker exec -e PGPASSWORD="$POSTGRES_PASSWORD" -i "$PG" \
  psql -U "$USER" -d "$DB" -v ON_ERROR_STOP=1 <<< "$SQL"

echo "== Counts after =="
docker exec -e PGPASSWORD="$POSTGRES_PASSWORD" "$PG" \
  psql -U "$USER" -d "$DB" -c "
SELECT relname AS table, n_live_tup AS approx_rows
FROM pg_stat_user_tables
ORDER BY relname;
"

echo "== Restart litellm (refresh Admin UI caches) =="
docker service update --force ai-router_litellm >/dev/null
for i in $(seq 1 40); do
  curl -sf --max-time 5 https://litellm.ift.calentiq.com/health/liveliness >/dev/null && break
  sleep 2
done
curl -sf --max-time 5 https://litellm.ift.calentiq.com/health/liveliness && echo
echo LITELLM_DB_CLEARED_OK
REMOTE