#!/bin/bash
#
# Copyright (C) 2026 Danilo M. <danix@danix.xyz>
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 2 as
# published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# gen-frames: make the companion's frames listed in frames.txt.
#
#   gen-frames [--seed N] [--list] [OUTPUT...]
#
# With no OUTPUT, makes every frame whose image does not exist yet. Naming
# outputs (dress/mad, leotard/mad) remakes those, existing or not; a dress
# frame remade leaves its leotard frame stale, so name both. --seed changes
# the seed (default 42) to reroll a result; --list prints what would be made.
# Then runs prep-frames.
#
# It drives a running sd-server rather than sd-cli, which would reload the
# model for every frame: start sd-webui and pick FLUX.2 klein first. sd-webui
# also unloads llama-server's model, which otherwise leaves too little VRAM.
# The model is checked from sd-server's command line, since its API does not
# name it.

set -euo pipefail

here=$(cd "$(dirname "$0")" && pwd)
art=${ART_DIR:-$HOME/.local/share/keystone/companion/2b}
url=${SD_URL:-http://127.0.0.1:7860}
seed=42
list=0
want=()
while (($#)); do
    case $1 in
        --seed) seed=$2; shift 2 ;;
        --list) list=1; shift ;;
        -h|--help) sed -n '14,28p' "$0" | sed 's/^# \{0,1\}//'; exit 0 ;;
        *) want+=("$1"); shift ;;
    esac
done

# frames.txt to "output<TAB>references<TAB>prompt" lines, variables expanded.
# References are joined with + and no spaces: the first is the image edited,
# any after it are extra images the prompt can point at ("the second image").
parse() {
    local line name out ref prompt
    local frame='^[[:space:]]*([a-z0-9/_-]+)[[:space:]]*<-[[:space:]]*([a-z0-9/_-]+([[:space:]]*\+[[:space:]]*[a-z0-9/_-]+)*)[[:space:]]*:[[:space:]]*(.*)$'
    declare -A vars=()
    while IFS= read -r line || [[ -n $line ]]; do
        [[ $line =~ ^[[:space:]]*(#|$) ]] && continue
        if [[ $line =~ ^([A-Z_]+)[[:space:]]*=[[:space:]]*(.*)$ ]]; then
            vars[${BASH_REMATCH[1]}]=${BASH_REMATCH[2]}
            continue
        fi
        if [[ ! $line =~ $frame ]]; then
            echo "gen-frames: frames.txt: cannot read: $line" >&2
            exit 1
        fi
        out=${BASH_REMATCH[1]} ref=${BASH_REMATCH[2]// /} prompt=${BASH_REMATCH[4]}
        for name in "${!vars[@]}"; do prompt=${prompt//\$\{$name\}/${vars[$name]}}; done
        if [[ $prompt == *'${'* ]]; then
            echo "gen-frames: frames.txt: $out: undefined variable in: $prompt" >&2
            exit 1
        fi
        printf '%s\t%s\t%s\n' "$out" "$ref" "$prompt"
    done <"$here/frames.txt"
}

# What to make: the named outputs, else every missing one, in file order.
# (No grep -q in a pipe: under pipefail its early exit fails the printf.)
todo=()
rows=$(parse) || exit 1
mapfile -t rows <<<"$rows"
declare -A known=() wanted=()
for row in "${rows[@]}"; do known[${row%%$'\t'*}]=1; done
for name in "${want[@]}"; do
    [[ -n ${known[$name]:-} ]] || { echo "gen-frames: $name is not in frames.txt" >&2; exit 1; }
    wanted[$name]=1
done
for row in "${rows[@]}"; do
    out=${row%%$'\t'*}
    if ((${#want[@]})); then
        [[ -n ${wanted[$out]:-} ]] && todo+=("$row")
    elif [[ ! -e $art/$out.png ]]; then
        todo+=("$row")
    fi
done
if ((list)); then
    for row in "${todo[@]}"; do printf '%s <- %s\n' "${row%%$'\t'*}" "$(cut -f2 <<<"$row")"; done
    exit 0
fi
if ((${#todo[@]} == 0)); then
    echo "gen-frames: nothing to make"
    exec "$here/prep-frames" "$art"
fi

pgrep -ax sd-server | grep -q klein ||
    { echo "gen-frames: sd-server is not running FLUX.2 klein; start sd-webui and pick it" >&2; exit 1; }

tmp=$(mktemp -d)
trap 'rm -rf "$tmp"' EXIT

made=0
for row in "${todo[@]}"; do
    IFS=$'\t' read -r out ref prompt <<<"$row"
    IFS=+ read -ra refs <<<"$ref"
    missing=""
    for r in "${refs[@]}"; do [[ -e $art/$r.png ]] || missing=$r; done
    if [[ -n $missing ]]; then
        echo "gen-frames: $out: its reference $missing.png does not exist yet, skipped" >&2
        continue
    fi
    echo "gen-frames: $out <- ${ref//+/ + }"
    # One JSON string per reference, as data URLs, slurped into the array.
    for r in "${refs[@]}"; do
        { printf 'data:image/png;base64,'; base64 -w0 "$art/$r.png"; } | jq -Rs .
    done >"$tmp/refs.json"
    # The settings the frames were tuned with in sd-webui: recipe.txt.
    jq -n --arg prompt "$prompt" --slurpfile refs "$tmp/refs.json" --argjson seed "$seed" '{
        prompt: $prompt, negative_prompt: "", width: 1024, height: 1024, seed: $seed,
        batch_count: 1, ref_images: $refs, auto_resize_ref_image: true,
        increase_ref_index: false, lora: [], output_format: "png", output_compression: 100,
        sample_params: { sample_steps: 4, sample_method: "euler", scheduler: "default",
                         guidance: { txt_cfg: 1.0 } }
    }' >"$tmp/body"
    id=$(curl -sf -m 30 -X POST "$url/sdcpp/v1/img_gen" -H 'Content-Type: application/json' \
        --data-binary @"$tmp/body" | jq -r '.id // empty') || id=""
    [[ -n $id ]] || { echo "gen-frames: $out: sd-server refused the job" >&2; exit 1; }
    while :; do
        curl -sf -m 30 "$url/sdcpp/v1/jobs/$id" >"$tmp/job" || { echo "gen-frames: $out: lost the job" >&2; exit 1; }
        status=$(jq -r '.status' "$tmp/job")
        case $status in
            queued|generating) sleep 2 ;;
            completed) break ;;
            *) echo "gen-frames: $out: $status: $(jq -r '.error.message // "no message"' "$tmp/job")" >&2; exit 1 ;;
        esac
    done
    mkdir -p "$art/$(dirname "$out")"
    jq -r '.result.images[0].b64_json' "$tmp/job" | base64 -d >"$tmp/out.png"
    mv "$tmp/out.png" "$art/$out.png"
    printf '%s  seed %s  %s <- %s\n' "$(date +%F)" "$seed" "$out" "$ref" >>"$art/generated.log"
    made=$((made + 1))
done
echo "gen-frames: made $made"
exec "$here/prep-frames" "$art"