#!/usr/bin/env bash
# Serve this model on TT hardware. Assumes ./install.sh has been run.
set -euo pipefail
HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
VENV="${VENV:-$HERE/venv}"
PYBIN="$VENV/bin/python"

# Locate ttnn WITHOUT importing it — importing loads _ttnn.so, which is exactly what needs the
# LD_PRELOAD below (chicken-and-egg). find_spec resolves the path without executing the module.
TTNN_DIR="$("$PYBIN" -c 'import importlib.util,os;print(os.path.dirname(importlib.util.find_spec("ttnn").origin))')"
export TT_METAL_RUNTIME_ROOT="$TTNN_DIR"
sfpi_meta="$TTNN_DIR/tt_metal/sfpi-version"
sfpi_version=""
sfpi_build=""
if [ -f "$sfpi_meta" ]; then
  sfpi_requirements="$("$PYBIN" -I - "$sfpi_meta" <<'PY'
import re
import sys
from pathlib import Path

values = {}
for line in Path(sys.argv[1]).read_text(encoding="utf-8").splitlines():
    match = re.fullmatch(r"(sfpi_(?:version|build))='([0-9.]+)'", line)
    if match:
        values[match.group(1)] = match.group(2)
if set(values) != {"sfpi_version", "sfpi_build"}:
    raise SystemExit("invalid or incomplete TTNN sfpi-version metadata")
print(values["sfpi_version"], values["sfpi_build"], sep="\t")
PY
)"
  IFS=$'\t' read -r sfpi_version sfpi_build <<<"$sfpi_requirements"
  sfpi_want="${sfpi_version}[${sfpi_build}]"
else
  # An older TTNN wheel declares no SFPI version: it is unknown, so skip that check.
  printf 'note: %s not found; skipping the SFPI version check\n' "$sfpi_meta" >&2
  sfpi_want="an unknown version"
fi
local_sfpi="$TTNN_DIR/runtime/sfpi"
system_sfpi=/opt/tenstorrent/sfpi
if [ -e "$local_sfpi/compiler/bin/riscv-tt-elf-g++" ]; then
  selected_sfpi="$local_sfpi"
elif [ -e "$system_sfpi/compiler/bin/riscv-tt-elf-g++" ]; then
  selected_sfpi="$system_sfpi"
else
  printf 'SFPI missing: TTNN requires %s; install the matching external host package\n' \
    "$sfpi_want" >&2
  exit 1
fi
gxx="$selected_sfpi/compiler/bin/riscv-tt-elf-g++"
test -x "$gxx" || { printf 'SFPI compiler is not executable: %s\n' "$gxx" >&2; exit 1; }
test -r "$selected_sfpi/include/sfpi_lib.h" || {
  printf 'SFPI headers are incomplete at %s\n' "$selected_sfpi" >&2
  exit 1
}
actual="$($gxx --version | sed -n '1p')"
[ -z "$sfpi_version" ] || case "$actual" in
  *"tenstorrent/sfpi:${sfpi_version}[${sfpi_build}]"*)
    printf 'validated SFPI %s[%s] at %s\n' "$sfpi_version" "$sfpi_build" "$selected_sfpi" >&2
    ;;
  *)
    # An OLDER SFPI lacks APIs the TTNN kernels use and fails at JIT compile, so it is fatal.
    found="$(printf '%s' "$actual" | sed -n 's/.*tenstorrent\/sfpi:\([0-9.]*\).*/\1/p')"
    if [ -n "$found" ] && [ "$(printf '%s\n%s\n' "$found" "$sfpi_version" | sort -V | head -1)" = "$found" ] \
        && [ "$found" != "$sfpi_version" ] && [ "${TT_MODEL_PRINT:-0}" != "1" ]; then
      printf 'SFPI at %s is older than TTNN needs (want %s[%s], found: %s); its kernels will not compile. Install SFPI %s or newer.\n' \
        "$selected_sfpi" "$sfpi_version" "$sfpi_build" "$actual" "$sfpi_version" >&2
      exit 1
    fi
    # A newer SFPI is usually compatible, so it only warns unless TT_MODEL_STRICT_SFPI=1.
    if [ "${TT_MODEL_STRICT_SFPI:-0}" = "1" ] && [ "${TT_MODEL_PRINT:-0}" != "1" ]; then
      printf 'SFPI mismatch at %s: need %s[%s], got: %s (TT_MODEL_STRICT_SFPI=1)\n' \
        "$selected_sfpi" "$sfpi_version" "$sfpi_build" "$actual" >&2
      exit 1
    fi
    printf 'warning: SFPI at %s is not the version TTNN declares (want %s[%s], found: %s); continuing. Set TT_MODEL_STRICT_SFPI=1 to make this fatal.\n' \
      "$selected_sfpi" "$sfpi_version" "$sfpi_build" "$actual" >&2
    ;;
esac

# _ttnncpp.so lives in ttnn.libs/ for an auditwheel-repaired (portable) wheel, or build/lib/ for a
# raw one; preload it to avoid the glibc "static TLS block" error on late dlopen.
# Prefer the auditwheel-vendored copy in *.libs/ (that's the one _ttnn.so actually loads via
# RPATH); fall back to build/lib for a raw (unrepaired) wheel — e.g. the plain `ttnn` PyPI wheel
# a v6 thin bundle installs, which isn't auditwheel-repaired at all.
# `|| true` on each probe: under `set -e -o pipefail`, a `ls <no-match> | head -1` pipe fails
# (pipefail surfaces ls's nonzero exit) and set -e would kill the script on THIS line, before the
# fallback below — or the deliberate `:?` error a few lines down — ever runs.
LD_PRELOAD="$(ls "$TTNN_DIR"/../*.libs/_ttnncpp*.so 2>/dev/null | head -1)" || true
[ -n "$LD_PRELOAD" ] || LD_PRELOAD="$(ls "$TTNN_DIR"/build/lib/_ttnncpp*.so 2>/dev/null | head -1)" || true
export LD_PRELOAD="${LD_PRELOAD:?could not locate _ttnncpp.so in the ttnn install}"
export TT_METAL_HOME="$TTNN_DIR"
# EXTRA_MODELS_DIR is a PARENT of per-model bundle folders; the plugin scans its children for
# each vllm_metadata.json (so the metadata lives in vllm_models/<model>/, not the bundle root).
export EXTRA_MODELS_DIR="$HERE/vllm_models"
export TT_VLLM_BUILTIN_MODELS=0
# Do NOT set VLLM_PLUGINS: it is an ALLOW-LIST — setting it suppresses the vllm.general_plugins
# group, so the model's tool/reasoning-parser overrides would silently not load. The TT platform
# + model registry load via entry points without it.
export PYTHONPATH="$HERE/extra_code:$HERE:${PYTHONPATH:-}"   # resolves the adapter/model imports
export MESH_DEVICE="${MESH_DEVICE:-P150}"
NCHIPS=1
if [ -z "${TT_METAL_VISIBLE_DEVICES:-}" ]; then
  if [ -n "${TT_VISIBLE_DEVICES:-}" ]; then
    IFS=, read -ra _GRANT <<< "$TT_VISIBLE_DEVICES"
    if [ "${#_GRANT[@]}" -lt "$NCHIPS" ]; then
      echo "run.sh: this model needs $NCHIPS chip(s) but TT_VISIBLE_DEVICES grants ${#_GRANT[@]} ($TT_VISIBLE_DEVICES)" >&2
      exit 1
    fi
    TT_VISIBLE_DEVICES="$(IFS=,; echo "${_GRANT[*]:0:$NCHIPS}")"
    export TT_VISIBLE_DEVICES
    TT_METAL_VISIBLE_DEVICES="0"
    _P150_MGD="$TTNN_DIR/tt_metal/fabric/mesh_graph_descriptors/p150_mesh_graph_descriptor.textproto"
    if [ -z "${TT_MESH_GRAPH_DESC_PATH:-}" ] && [ -f "$_P150_MGD" ]; then
      export TT_MESH_GRAPH_DESC_PATH="$_P150_MGD"
    fi
  else
    TT_METAL_VISIBLE_DEVICES="0"
  fi
fi
export TT_METAL_VISIBLE_DEVICES
# HERMETIC RUNTIME: keep every cache/home INSIDE the folder wall, so serving writes and reads
# nothing outside it (the ttnn tensor cache even DEFAULTS to a hard-coded /mnt/... path upstream —
# a classic other-machine leak we must override). Each is overridable if the operator sets it.
export HF_HOME="${HF_HOME:-$HERE/.hf}"                  # HF weights + hub cache
export TT_CACHE_PATH="${TT_CACHE_PATH:-$HERE/.tt_cache}"    # ttnn weight/tensor cache
export TT_CACHE_HOME="${TT_CACHE_HOME:-$HERE/.tt_cache}"    # override upstream's /mnt/... default
export XDG_CACHE_HOME="${XDG_CACHE_HOME:-$HERE/.cache}"     # generic catch-all (triton, etc.)
export TT_METAL_CACHE="${TT_METAL_CACHE:-$HERE/.cache}"     # compiled TT-Metal kernels
export TRITON_CACHE_DIR="${TRITON_CACHE_DIR:-$HERE/.cache/triton}"
export TORCHINDUCTOR_CACHE_DIR="${TORCHINDUCTOR_CACHE_DIR:-$HERE/.cache/inductor}"
# ttnn.CONFIG defaults to $HOME/.cache/ttnn and reads only TTNN_CONFIG_OVERRIDES (not XDG).
if [ -z "${TTNN_CONFIG_OVERRIDES:-}" ]; then
  export TTNN_CONFIG_OVERRIDES="{\"cache_path\": \"$HERE/.cache/ttnn\", \"model_cache_path\": \"$HERE/.cache/ttnn/models\"}"
fi
export HF_MODEL="${HF_MODEL:-episod/minimax-music3-fiona-crapple}"
export TT_MODEL_WEIGHTS_REVISION="${TT_MODEL_WEIGHTS_REVISION:-53fc7557ea3101d3efbe5e50598bd0bea7f44344}"
export MUSIC3_FABRIC_CONFIG="none"
export MUSIC3_MESH_SHAPE="1x1"
export MUSIC3_MAX_SEQ_LEN="16384"
export MUSIC3_MAX_FRAMES="9000"
export MUSIC3_WARMUP="1"
export MUSIC3_WARMUP_FRAMES="200"
export MUSIC3_TRACE_REGION_MB="256"
export TT_METAL_OPERATION_TIMEOUT_SECONDS="180.0"
CMD=("$PYBIN" -m uvicorn --host "${HOST:-127.0.0.1}" --port "${PORT:-20150}" --lifespan on fiona_asgi:app "$@")
# TT_MODEL_PRINT=1 (set by `tt-model serve --print`) echoes the fully-resolved command+env
if [ "${TT_MODEL_PRINT:-0}" = "1" ]; then
  printf 'LD_PRELOAD=%s TT_METAL_HOME=%s EXTRA_MODELS_DIR=%s MESH_DEVICE=%s HF_MODEL=%s
  ' \
    "$LD_PRELOAD" "$TT_METAL_HOME" "$EXTRA_MODELS_DIR" "$MESH_DEVICE" "${HF_MODEL:-}"
  printf '%q ' "${CMD[@]}"; echo  # shell-quoted so the printed command pastes back intact
  exit 0
fi
exec "${CMD[@]}"