#!/usr/bin/env bash
# TOASTSIM asset generation — 100% local, 100% free, via MODELBEAST on the M3 Ultra.
#
#   ./scripts/gen-assets.sh 3d     # product shot -> cutout -> GLB  (slow: ~5 min each, gpu lane is serial)
#   ./scripts/gen-assets.sh 2d     # flux images (judge portraits, backdrops, ui)
#   ./scripts/gen-assets.sh all
#
# Seeds are fixed so every asset is re-generatable. Style token is shared for coherence.
set -uo pipefail

MB_DIR="$HOME/Documents/MODELBEAST"
REPO="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
OUT_MODELS="$REPO/public/assets/models"
OUT_IMG="$REPO/public/assets/img"
mkdir -p "$OUT_MODELS" "$OUT_IMG"

cd "$MB_DIR" || exit 1
# shellcheck disable=SC1091
source data/agent.env

STYLE="chunky stylized claymation kitchen prop, soft matte plastic look, warm colors"
SHOT="single object centered on plain light grey background, soft studio lighting, 3/4 view, chunky stylized game asset"

# run_flux <seed> <prompt> [w] [h] -> asset id on stdout
run_flux() {
  local seed="$1" prompt="$2" w="${3:-1024}" h="${4:-1024}"
  ./mb run flux_local -p prompt="$prompt" -p model=flux2-klein-4b -p steps=4 \
      -p seed="$seed" -p width="$w" -p height="$h" --wait 2>/dev/null \
    | awk '/image/{print $1}' | tail -1
}

# gen_3d <name> <seed> <prompt>
gen_3d() {
  local name="$1" seed="$2" prompt="$3"
  if [ -f "$OUT_MODELS/$name.glb" ]; then echo "  = $name.glb exists, skipping"; return; fi
  echo "→ [$name] flux..."
  local img cut
  img=$(run_flux "$seed" "$prompt, $STYLE, $SHOT")
  [ -z "$img" ] && { echo "  ! $name: flux failed"; return 1; }
  echo "→ [$name] cutout ($img)..."
  cut=$(./mb run bg_remove_local --asset "$img" -p resolution=2048 --wait 2>/dev/null | awk '/image/{print $1}' | tail -1)
  [ -z "$cut" ] && { echo "  ! $name: bg_remove failed"; return 1; }
  echo "→ [$name] mesh ($cut) — ~5 min..."
  local tmp; tmp=$(mktemp -d)
  ./mb run hunyuan3d_mlx --asset "$cut" -p octree_resolution=256 -p texture_size=1024 \
      -p max_num_view=6 --wait --download "$tmp" >/dev/null 2>&1
  local glb; glb=$(find "$tmp" -iname '*.glb' | head -1)
  if [ -n "$glb" ]; then mv "$glb" "$OUT_MODELS/$name.glb"; echo "  ✓ $name.glb"; else echo "  ! $name: no glb"; fi
  rm -rf "$tmp"
}

# gen_2d <name> <seed> <prompt> [w] [h] [cutout]
gen_2d() {
  local name="$1" seed="$2" prompt="$3" w="${4:-1024}" h="${5:-1024}" cutout="${6:-no}"
  if [ -f "$OUT_IMG/$name.png" ]; then echo "  = $name.png exists, skipping"; return; fi
  echo "→ [$name] flux..."
  local img; img=$(run_flux "$seed" "$prompt" "$w" "$h")
  [ -z "$img" ] && { echo "  ! $name: flux failed"; return 1; }
  if [ "$cutout" = "cut" ]; then
    echo "→ [$name] cutout..."
    local c; c=$(./mb run bg_remove_local --asset "$img" -p resolution=2048 --wait 2>/dev/null | awk '/image/{print $1}' | tail -1)
    [ -n "$c" ] && img="$c"
  fi
  local tmp; tmp=$(mktemp -d)
  ./mb get "$img" -o "$tmp" >/dev/null 2>&1
  local f; f=$(find "$tmp" -iname '*.png' -o -iname '*.jpg' | head -1)
  [ -n "$f" ] && { mv "$f" "$OUT_IMG/$name.png"; echo "  ✓ $name.png"; }
  rm -rf "$tmp"
}

# ---- the judge: same person every time, seed fixed, only the expression phrase varies ----
JUDGE="portrait of a stern elderly australian food inspector, deep wrinkles, bushy grey eyebrows, thin wire glasses, grey moustache, wearing a brown cardigan over a shirt and tie, holding a clipboard, plain light grey background, head and shoulders, claymation stop-motion character, soft matte clay texture, warm studio light"

do_3d() {
  gen_3d toaster       11 "a retro two-slot pop-up toaster, cream and chrome, with a browning dial"
  gen_3d butter_dish   23 "a ceramic butter dish with a thick block of yellow butter on it, lid off"
  gen_3d pb_jar        31 "an open glass jar of peanut butter, chunky brown spread visible, cream label"
  gen_3d mitey_jar     37 "a small squat glass jar of dark brown yeast extract spread with a bright red and yellow label reading MITEY, lid off"
  gen_3d plate         43 "a simple round white ceramic side plate, empty, seen at an angle"
}

do_2d() {
  gen_2d judge_neutral      101 "$JUDGE, neutral unimpressed expression" 768 768 cut
  gen_2d judge_intrigued    101 "$JUDGE, one eyebrow raised, mildly intrigued expression" 768 768 cut
  gen_2d judge_impressed    101 "$JUDGE, small approving smile, impressed expression" 768 768 cut
  gen_2d judge_disappointed 101 "$JUDGE, deeply disappointed frowning expression, looking down at clipboard" 768 768 cut
  gen_2d judge_horrified    101 "$JUDGE, horrified appalled expression, eyes wide, recoiling" 768 768 cut
  gen_2d bench_wood         202 "seamless top-down texture of a warm honey-toned wooden kitchen benchtop, subtle grain and knife marks, claymation soft matte look, flat even lighting, no objects" 1024 1024
  gen_2d kitchen_backdrop   205 "blurry cosy 1970s australian kitchen interior background, warm cream tiles, soft bokeh, claymation set, no people, no text" 1280 768
  gen_2d ticket_paper       210 "seamless texture of a plain cream paper order docket, faint blue lines, slightly crumpled, flat even lighting, no text" 512 512
  gen_2d title_art          220 "a single slice of golden toast with a pat of butter melting on it, dramatic hero shot, claymation stop-motion, soft matte clay, warm rim light, plain dark background" 1024 768 cut
}

case "${1:-all}" in
  3d) do_3d ;;
  2d) do_2d ;;
  *)  do_3d; do_2d ;;
esac
echo "done."