#!/bin/bash
set -euo pipefail

# ===== Paths =====
BODY="${BODY:-/Users/josephanady/Desktop/AI_Body}"
BRAIN="${BRAIN:-/Users/josephanady/Desktop/AI_Brain_Control}"
DB="$BRAIN/db/shards_index.sqlite"
LOG_DIR="$BRAIN/logs"
DOC="$BRAIN/docs/body_brain_builder.md"

# ===== Config =====
PROFILE="${PROFILE:-ssd}"                 # nvme|ssd|hdd
TOTAL_SHARDS="${TOTAL_SHARDS:-0}"         # 0 = run forever
TARGET_COMPRESSED_MB="${TARGET_COMPRESSED_MB:-512}"
COMPRESSOR="${COMPRESSOR:-auto}"          # auto|zstd|gzip
COMPRESS_LEVEL="${COMPRESS_LEVEL:-3}"
CAFFEINATE="${CAFFEINATE:-1}"

SHARDS_PER_BATCH="${SHARDS_PER_BATCH:-12}"
SLEEP_BETWEEN_SHARDS="${SLEEP_BETWEEN_SHARDS:-1}"
SLEEP_BETWEEN_BATCHES="${SLEEP_BETWEEN_BATCHES:-15}"

# DB safety knobs
BUSY_TIMEOUT_MS="${BUSY_TIMEOUT_MS:-5000}"      # PRAGMA busy_timeout / .timeout
MAX_DB_RETRIES="${MAX_DB_RETRIES:-12}"          # INSERT retries on lock
RETRY_BASE_MS="${RETRY_BASE_MS:-80}"            # base backoff (ms), with jitter

# Concurrency knobs
MAX_CONCURRENCY="${MAX_CONCURRENCY:-$SHARDS_PER_BATCH}"  # cap parallel shard jobs
PRE_INSERT_JITTER_MS="${PRE_INSERT_JITTER_MS:-50}"        # 0..N ms random jitter before DB insert

case "$PROFILE" in
  nvme) : "${TARGET_COMPRESSED_MB:=1024}"; : "${COMPRESS_LEVEL:=3}"; : "${SHARDS_PER_BATCH:=30}"; : "${SLEEP_BETWEEN_SHARDS:=0}";  : "${SLEEP_BETWEEN_BATCHES:=10}";;
  ssd)  : "${TARGET_COMPRESSED_MB:=512}";  : "${COMPRESS_LEVEL:=3}"; : "${SHARDS_PER_BATCH:=20}"; : "${SLEEP_BETWEEN_SHARDS:=1}";  : "${SLEEP_BETWEEN_BATCHES:=20}";;
  hdd)  : "${TARGET_COMPRESSED_MB:=384}";  : "${COMPRESS_LEVEL:=1}"; : "${SHARDS_PER_BATCH:=8}";  : "${SLEEP_BETWEEN_SHARDS:=2}";  : "${SLEEP_BETWEEN_BATCHES:=60}";;
esac

# Tools
if [ "$COMPRESSOR" = "auto" ]; then
  if command -v zstd >/dev/null 2>&1; then COMPRESSOR="zstd"
  elif command -v gzip >/dev/null 2>&1; then COMPRESSOR="gzip"
  else echo "No zstd or gzip found."; exit 1; fi
fi
EXT="zst"; [ "$COMPRESSOR" = "gzip" ] && EXT="gz"

# caffeinate wrapper
if [ "$CAFFEINATE" -eq 1 ] && command -v caffeinate >/dev/null; then
  exec caffeinate -dimsu "$0" "$@"
fi

# Ensure DB ready
"$BRAIN/scripts/init_db.sh" "$DB" >/dev/null 2>&1 || true
source "$BRAIN/scripts/_util.sh"

mkdir -p "$BODY" "$LOG_DIR"

# ---- helpers ----
sql_escape() {  # escape single quotes for SQL literal
  printf "%s" "$1" | sed "s/'/''/g"
}

ms_sleep() {  # sleep in milliseconds with ms precision
  local ms="$1"
  if command -v python3 >/dev/null 2>&1; then
    python3 - "$ms" <<'PY'
import time, sys
time.sleep(float(sys.argv[1]) / 1000.0)
PY
  else
    sleep "$(awk -v m="$ms" 'BEGIN{printf "%.3f", m/1000}')"
  fi
}

sqlite_insert_retry() {
  # Args: SQL
  local sql="$1"
  local n=0
  while (( n <= MAX_DB_RETRIES )); do
    if echo ".timeout $BUSY_TIMEOUT_MS; $sql" | sqlite3 "$DB" ; then
      return 0
    fi
    n=$((n+1))
    # backoff with jitter: base * (1+n) +/- 0..base
    jitter=$(( RANDOM % RETRY_BASE_MS ))
    delay=$(( (RETRY_BASE_MS*(1+n)) + jitter ))
    echo "DB busy; retry $n/$MAX_DB_RETRIES after ${delay}ms"
    ms_sleep "$delay"
  done
  echo "INSERT failed after $MAX_DB_RETRIES retries"
  return 1
}

sqlite_select_safe() {
  # SELECT with busy_timeout for telemetry/reads
  echo ".timeout $BUSY_TIMEOUT_MS; $1" | sqlite3 "$DB"
}

# Concurrency helpers
running_jobs() {
  jobs -rp | wc -l | tr -d ' '
}

await_capacity() {
  local max=${1:-$MAX_CONCURRENCY}
  while [ "$(running_jobs)" -ge "$max" ]; do
    sleep 0.05
  done
}

next_id() {
  # prefer DB, fallback to files
  if command -v sqlite3 >/dev/null 2>&1; then
    local last
    last=$(sqlite_select_safe "SELECT COALESCE(MAX(id),-1) FROM shards;")
    echo $((last + 1))
  else
    local n
    n=$(ls "$BODY"/shard_*.${EXT} 2>/dev/null | sed -E 's/.*shard_([0-9]+)\.'${EXT}'/\1/' | sort -n | tail -1 || echo "")
    [ -z "$n" ] && echo 0 || echo $((n+1))
  fi
}

compress_one() {
  local id="$1"
  local shard="$BODY/shard_${id}.${EXT}"
  local shard_esc started completed bytes hash thru
  shard_esc="$(sql_escape "$shard")"
  started="$(date -u +%FT%TZ)"

  echo ">>> [$started] Building shard #$id (${TARGET_COMPRESSED_MB}MB, $COMPRESSOR -$COMPRESS_LEVEL) → $shard"
  local t0=$(python3 - <<'PY'
import time; print(time.time())
PY
)
  if [ "$COMPRESSOR" = "zstd" ]; then
    dd if=/dev/urandom bs=1m count="$TARGET_COMPRESSED_MB" 2>/dev/null | zstd -"${COMPRESS_LEVEL}" -q -o "$shard"
  else
    dd if=/dev/urandom bs=1m count="$TARGET_COMPRESSED_MB" 2>/dev/null | gzip -"${COMPRESS_LEVEL}" -c > "$shard"
  fi

  bytes=$(os_file_size_bytes "$shard")
  hash=$(sha256_file "$shard")

  local t1=$(python3 - <<'PY'
import time; print(time.time())
PY
)
  local dur=$(python3 - <<PY
t0=$t0; t1=$t1
print(max(0.001, t1 - t0))
PY
)
  thru=$(python3 - <<PY
print(round(($bytes/1048576)/$dur, 3))
PY
)
  completed="$(date -u +%FT%TZ)"

  # Optional tiny jitter to decorrelate concurrent INSERTs
  if [ "$PRE_INSERT_JITTER_MS" -gt 0 ]; then
    ms_sleep $(( RANDOM % (PRE_INSERT_JITTER_MS + 1) ))
  fi

  # INSERT with timeout + retry/backoff
  sqlite_insert_retry "
INSERT INTO shards(id,file,size_bytes,hash,started_at,completed_at,profile,compressor,throughput_mb_s)
VALUES($id,'$shard_esc',$bytes,'$hash','$started','$completed','$PROFILE','$COMPRESSOR',$thru)
ON CONFLICT(id) DO UPDATE SET
  file=excluded.file,
  size_bytes=excluded.size_bytes,
  hash=excluded.hash,
  started_at=excluded.started_at,
  completed_at=excluded.completed_at,
  profile=excluded.profile,
  compressor=excluded.compressor,
  throughput_mb_s=excluded.throughput_mb_s;
" || { echo "❌ failed to index shard #$id"; return 1; }

  echo "✅ [$completed] Shard #$id done ($((bytes/1048576)) MB, ${thru} MB/s)"
}

refresh_doc() {
  mkdir -p "$(dirname "$DOC")"
  {
    echo "# AI Body/Brain Builder – Live Handshake"
    echo
    echo "Last Updated (UTC): $(date -u +%FT%TZ)"
    echo
    echo "## DB safety"
    echo "- busy_timeout: ${BUSY_TIMEOUT_MS} ms"
    echo "- retries: ${MAX_DB_RETRIES} (base ${RETRY_BASE_MS}ms + jitter)"
    echo
    echo "## Paths"
    echo "- BODY:  $BODY"
    echo "- BRAIN: $BRAIN"
    echo "- DB:    $DB"
    echo
    echo "## Profile"
    echo "- profile: $PROFILE"
    echo "- compressor: $COMPRESSOR (-$COMPRESS_LEVEL)"
    echo "- target shard size: ${TARGET_COMPRESSED_MB} MB"
    echo "- concurrency: up to ${MAX_CONCURRENCY} parallel"
    echo
    echo "## Telemetry (last 10)"
    if command -v sqlite3 >/dev/null 2>&1; then
      sqlite_select_safe "SELECT id,size_bytes,throughput_mb_s,completed_at FROM shards ORDER BY id DESC LIMIT 10;" | sed 's/|/  |  /g'
    else
      echo "(sqlite3 not found; telemetry table unavailable)"
    fi
    echo
    echo "## How to run"
    echo "\\`TOTAL_SHARDS=2 PROFILE=ssd TARGET_COMPRESSED_MB=128 $0\\`"
  } > "$DOC"
  echo "🤝 Handshake refreshed → $DOC"
}

main() {
  local id built
  id=$(next_id)
  echo "Resuming at shard id = $id"
  built=0
  while [ "$TOTAL_SHARDS" -eq 0 ] || [ $built -lt "$TOTAL_SHARDS" ]; do
    for _ in $(seq 1 "$SHARDS_PER_BATCH"); do
      if [ "$TOTAL_SHARDS" -ne 0 ] && [ $built -ge "$TOTAL_SHARDS" ]; then break; fi
      await_capacity "$MAX_CONCURRENCY"
      compress_one "$id" &
      sleep "$SLEEP_BETWEEN_SHARDS"
      id=$((id+1))
      built=$((built+1))
    done
    wait
    refresh_doc
    echo "--- Batch complete, sleeping ${SLEEP_BETWEEN_BATCHES}s ---"
    sleep "$SLEEP_BETWEEN_BATCHES"
  done
  echo "All done. Total new shards built: $built"
}

main "$@"