#!/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"
[ -d "$MB_DIR" ] || MB_DIR="$HOME/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
if [ -f data/agent.env ]; then source data/agent.env
else set -a; source "$HOME/Documents/backnforth/.env"; set +a; fi
export MB_HOST="${MB_HOST:-http://100.89.131.57:8777}"

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"

# ponytail: m4pro node has dead venvs (2026-07-17) — jobs routed there fail in ~1s,
# so every stage retries up to 10x with a 60s pause until a healthy node takes it.
# Remove the loops once m4pro is fixed or pulled from the pool.
MB_TRIES=10 MB_PAUSE=60

# run_flux <seed> <prompt> [w] [h] -> asset id on stdout
run_flux() {
  local seed="$1" prompt="$2" w="${3:-1024}" h="${4:-1024}" i out
  for i in $(seq 1 "$MB_TRIES"); do
    out=$(./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)
    [ -n "$out" ] && { echo "$out"; return; }
    sleep "$MB_PAUSE"
  done
}

# run_cutout <asset_id> -> asset id on stdout
run_cutout() {
  local asset="$1" i out
  for i in $(seq 1 "$MB_TRIES"); do
    out=$(./mb run bg_remove_local --asset "$asset" -p resolution=2048 --wait 2>/dev/null \
      | awk '/image/{print $1}' | tail -1)
    [ -n "$out" ] && { echo "$out"; return; }
    sleep "$MB_PAUSE"
  done
}

# gen_3d <name> <seed> <prompt> — each stage retries once (transient rsync/infra errors happen)
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")
  if [ -z "$img" ]; then
    echo "  ~ $name: flux failed, retrying once (seed+1)..."
    img=$(run_flux "$((seed+1))" "$prompt, $STYLE, $SHOT")
  fi
  [ -z "$img" ] && { echo "  ! $name: flux failed twice"; return 1; }
  echo "→ [$name] cutout ($img)..."
  cut=$(run_cutout "$img")
  [ -z "$cut" ] && { echo "  ! $name: bg_remove failed after $MB_TRIES tries"; return 1; }
  echo "→ [$name] mesh ($cut) — ~5 min..."
  local try glb
  for try in $(seq 1 "$MB_TRIES"); do
    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
    glb=$(find "$tmp" -iname '*.glb' | head -1)
    if [ -n "$glb" ]; then mv "$glb" "$OUT_MODELS/$name.glb"; echo "  ✓ $name.glb"; rm -rf "$tmp"; return 0; fi
    rm -rf "$tmp"
    echo "  ~ $name: no glb (try $try/$MB_TRIES), pausing ${MB_PAUSE}s..."
    sleep "$MB_PAUSE"
  done
  echo "  ! $name: no glb after $MB_TRIES tries"
  return 1
}

# 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=$(run_cutout "$img")
    [ -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
}

# ---- M10: the bakery (loaves), plus bruschetta ingredients for the milestone after ----
# (seeds bumped +1 and prompts reworded after the first batch died to transient rsync errors)
do_bakery() {
  gen_3d sourdough_loaf 62 "a round rustic sourdough boule of bread, thick crackled golden crust, flour dusted, scored cross on top, whole and unsliced"
  gen_3d batard_loaf    68 "an oval batard loaf of crusty artisan bread, golden brown crust with diagonal score marks, whole and unsliced"
  gen_3d cutting_board  72 "a thick rectangular wooden chopping board, rounded corners, warm honey wood with worn grain"
  gen_3d tomato_vine    74 "three plump ripe red tomatoes attached to one green vine stem, glossy skin"
  gen_3d roast_tomatoes 80 "a small metal roasting tray holding blistered oven-roasted tomato halves, charred edges, glistening"
  gen_3d cheese_hard    83 "a wedge of hard aged parmesan style cheese, pale gold, cracked granular texture"
  gen_3d cheese_soft    89 "a small round of soft white brie style cheese with one wedge cut out showing creamy interior"
}

# ---- M12–M15: the prep bench cast (Build Brief 2 §8) ----
do_prep() {
  gen_3d orange_whole    103 "a single whole orange with dimpled peel and a small green stem nub"
  gen_3d orange_half     107 "half an orange cut face up, showing juicy segments in a radial pattern and white pith ring"
  gen_3d juicer_pyramid  109 "an old fashioned glass citrus juicer reamer with a tall ribbed pyramid cone in the centre and a moat rim, pressed glass ribs clearly visible"
  gen_3d juicer_deluxe   113 "a sturdy pressed glass citrus juicer with a ribbed reaming cone, deep seed-catching moat with strainer notches, art deco look"
  gen_3d juicernaut      127 "a magnificent heavy chromed citrus juicer with a fluted ribbed cone, gleaming polished metal, trophy-like, over-engineered"
  gen_3d juice_glass     131 "a small plain drinking glass tumbler, empty, slightly tapered"
  gen_3d box_grater      137 "a stainless steel box grater with four faces and a rolled handle on top"
  gen_3d onion_brown     139 "a whole brown onion with papery golden skin and a wispy root"
  gen_3d garlic_bulb     149 "a whole garlic bulb, papery white skin with purple streaks"
  gen_3d basil_sprig     151 "a small sprig of fresh basil with several green leaves on a stem"
  gen_3d oil_bottle      157 "a tall glass olive oil bottle with a narrow spout, golden green oil visible inside"
  gen_3d oven_tray       163 "a plain rectangular metal oven baking tray with a rolled edge"
}

# ---- The pantry backlog cast (brief §7.5): mushrooms + bell pepper ----
do_pantry() {
  gen_3d button_mushroom 167 "a single small white button mushroom, smooth rounded cap, short stubby stem"
  gen_3d portobello      171 "a large flat portobello mushroom cap face down showing brown gills, thick pale stem"
  gen_3d king_oyster     173 "a king oyster mushroom, one thick smooth cream stem with a small light brown cap"
  gen_3d bell_pepper     179 "a whole glossy red bell pepper with a short green stem, deep vertical lobes"
}

# ---- FIRE & ENVIRONMENTS: heat sources, pans, outdoor, utensils (seeds 181-219) ----
# The heat-source cast. Universal props for the gas/electric/induction/charcoal
# + indoor/outdoor doctrine. Fire freely — all regenerable, all free.
do_heat() {
  # cooktops / heat sources
  gen_3d gas_hob        181 "a retro four burner gas stovetop cooktop, black cast iron trivets over blue gas burner rings, cream enamel body, control knobs at the front"
  gen_3d electric_hob   185 "a retro electric coil cooktop with four spiral metal heating elements on a cream enamel top, one element glowing faint orange"
  gen_3d induction_hob  189 "a sleek modern flat black glass induction cooktop, four faint printed ring outlines, minimal, glossy"
  gen_3d camp_stove     193 "a small portable single burner camping gas stove, folding pot supports, a green gas canister screwed to the side"
  # pans and vessels
  gen_3d cast_iron_pan  197 "a heavy black cast iron skillet frying pan with a long handle and a small pour spout, seasoned matte surface"
  gen_3d frying_pan     201 "a stainless steel frying pan with a riveted handle, shallow sloped sides, empty"
  gen_3d wok            205 "a round bottomed carbon steel wok with a single long wooden handle, dark seasoned interior"
  gen_3d saucepan       207 "a small stainless steel saucepan with a lid and a single handle"
  # outdoor
  gen_3d kettle_bbq     209 "a round black charcoal kettle barbecue with a domed vented lid and three slim legs, classic backyard bbq"
  gen_3d gas_bbq        211 "a rectangular outdoor gas barbecue grill with a lift up hood, side shelves, and a gas bottle underneath"
  # utensils
  gen_3d spatula        213 "a flat slotted metal turner spatula fish slice with a black handle"
  gen_3d wooden_spoon   215 "a plain wooden cooking spoon, worn honey timber"
  gen_3d tongs          217 "a pair of stainless steel kitchen tongs, spring hinge, scalloped gripping ends"
  gen_3d ladle          219 "a stainless steel soup ladle with a deep round bowl and a long hooked handle"
}

do_prep2d() {
  gen_2d bruschetta_hero 230 "rustic bruschetta on toasted sourdough, roasted tomato pieces and basil on top, dramatic hero shot, claymation stop-motion, soft matte clay, warm rim light, plain dark background" 1024 768 cut
  gen_2d wonderslice_bag 233 "a cheerful retro plastic bread bag with bold cartoon lettering style packaging design, red and yellow, claymation prop look, plain light grey background, no readable text" 768 768 cut
}

case "${1:-all}" in
  3d) do_3d ;;
  2d) do_2d ;;
  bakery) do_bakery ;;
  prep) do_prep; do_prep2d ;;
  pantry) do_pantry ;;
  heat) do_heat ;;
  *)  do_3d; do_2d ;;
esac
echo "done."