#!/bin/bash
# Copyright (C) 2026 Danilo M. <danix@danix.xyz>
# Licensed under the GNU General Public License version 2.
#
# Host-side wrapper for local SDXL image generation on Intel Arc.
# Subcommands manage a persistent daemon (model resident in VRAM); a bare
# prompt talks to the daemon if up, else falls back to a one-shot run.
#
# Usage:
#   imggen start [-m realvis]     # spin up daemon, load model, hold VRAM
#   imggen stop                   # kill daemon, free VRAM
#   imggen status                 # is it up? which model?
#   imggen list                   # show model keys
#   imggen "a red bicycle" [-m realvis] [--steps 30] [--seed 42] [-o out.png] [-n "blurry"]

set -euo pipefail

IMGGEN_DIR="${IMGGEN_DIR:-$HOME/Programming/GIT/imggen}"
cd "$IMGGEN_DIR"

# Load .env for OUT dir + port on the host side.
if [ -f .env ]; then set -a; . ./.env; set +a; fi

OUT_DIR="${IMGGEN_OUT:-$HOME/Pictures/AI-gen}"
PORT="${IMGGEN_PORT:-8765}"
BASE="http://127.0.0.1:${PORT}"
mkdir -p "$OUT_DIR"

usage() {
    echo "usage:" >&2
    echo "  imggen start [-m KEY] | stop | status | list" >&2
    echo "  imggen \"<prompt>\" [-m KEY] [--steps N] [--guidance G] [--seed S] [-n NEG] [-o name.png]" >&2
    exit 1
}

daemon_up() { curl -sf "${BASE}/status" >/dev/null 2>&1; }

case "${1:-}" in
    start)
        shift
        model="${IMGGEN_MODEL:-sdxl}"
        if [ "${1:-}" = "-m" ] || [ "${1:-}" = "--model" ]; then model="$2"; fi
        IMGGEN_MODEL="$model" docker compose up -d
        echo "daemon starting with model '$model' (VRAM ~9GB held until 'imggen stop')"
        exit 0
        ;;
    stop)
        docker compose down
        echo "daemon stopped, VRAM freed"
        exit 0
        ;;
    status)
        if daemon_up; then curl -s "${BASE}/status"; echo; else echo "daemon down"; fi
        exit 0
        ;;
    list)
        echo "sdxl     general (default)"
        echo "realvis  photoreal, permissive"
        echo "pony     unrestricted, versatile"
        exit 0
        ;;
    "" ) usage ;;
    -h|--help) usage ;;
esac

# --- prompt path ---
PROMPT="$1"; shift
MODEL=""; NEG=""; STEPS=""; GUID=""; SEED=""; OUT=""
W="${IMGGEN_WIDTH:-}"; H="${IMGGEN_HEIGHT:-}"
while [ "$#" -gt 0 ]; do
    case "$1" in
        -m|--model)    MODEL="$2"; shift 2 ;;
        -n|--negative) NEG="$2"; shift 2 ;;
        -s|--steps)    STEPS="$2"; shift 2 ;;
        -g|--guidance) GUID="$2"; shift 2 ;;
        --seed)        SEED="$2"; shift 2 ;;
        -W|--width)    W="$2"; shift 2 ;;
        -H|--height)   H="$2"; shift 2 ;;
        -o|--out)      OUT="$2"; shift 2 ;;
        *) echo "unknown arg: $1" >&2; usage ;;
    esac
done

if daemon_up; then
    # No -m: use the resident model as-is (no menu, no swap). The user already
    # chose it at 'imggen start'; -m KEY is the explicit override that swaps.
    MODEL="${MODEL:-$(curl -s "${BASE}/status" | sed -n 's/.*"model": *"\([^"]*\)".*/\1/p')}"
    MODEL="${MODEL:-sdxl}"
    OUT_NAME="${OUT:-${MODEL}-$(date +%Y%m%d-%H%M%S).png}"
    OUT_PATH="${OUT_DIR}/${OUT_NAME}"

    # Swap resident model if needed.
    current="$(curl -s "${BASE}/status" | sed -n 's/.*"model": *"\([^"]*\)".*/\1/p')"
    if [ "$current" != "$MODEL" ]; then
        echo "swapping daemon model $current -> $MODEL ..." >&2
        # Capture body and status; a failed swap (bad/gated repo, OOM) must be
        # surfaced, not silently swallowed by set -e killing the script.
        swap_body="$(curl -s -w '\n%{http_code}' -X POST "${BASE}/model" \
            -H 'Content-Type: application/json' -d "{\"name\":\"$MODEL\"}")"
        swap_code="${swap_body##*$'\n'}"
        swap_msg="${swap_body%$'\n'*}"
        if [ "$swap_code" != "200" ]; then
            echo "swap to '$MODEL' failed (HTTP $swap_code): $swap_msg" >&2
            exit 1
        fi
    fi

    # Build JSON request. jq not assumed; hand-build (values are ours/simple).
    req="{\"prompt\":$(printf '%s' "$PROMPT" | python3 -c 'import json,sys; print(json.dumps(sys.stdin.read()))')"
    [ -n "$NEG" ]   && req="$req,\"negative\":$(printf '%s' "$NEG" | python3 -c 'import json,sys; print(json.dumps(sys.stdin.read()))')"
    [ -n "$STEPS" ] && req="$req,\"steps\":$STEPS"
    [ -n "$GUID" ]  && req="$req,\"guidance\":$GUID"
    [ -n "$SEED" ]  && req="$req,\"seed\":$SEED"
    [ -n "$W" ]     && req="$req,\"width\":$W"
    [ -n "$H" ]     && req="$req,\"height\":$H"
    req="$req}"

    # No -f: keep the error body so we can classify it (e.g. GPU wedged).
    code="$(curl -s -w '%{http_code}' -X POST "${BASE}/generate" \
        -H 'Content-Type: application/json' -d "$req" -o "$OUT_PATH")"
    if [ "$code" = 200 ]; then
        echo "Saved ${OUT_PATH}"
    else
        # On failure the body is the server's JSON error, not a PNG.
        err="$(cat "$OUT_PATH" 2>/dev/null)"
        rm -f "$OUT_PATH"
        # level_zero DEVICE_LOST = GuC/xe driver hang; only a reboot recovers.
        if printf '%s' "$err" | grep -qi "device.lost\|DEVICE_LOST\|error 20"; then
            echo "GPU wedged (level_zero DEVICE_LOST): GuC/xe driver fault, reboot needed" >&2
        else
            echo "generate failed: ${err:-see daemon logs: docker compose logs}" >&2
        fi
        exit 1
    fi
else
    # One-shot fallback: non-interactive, model defaults to sdxl or -m KEY.
    echo "daemon down, one-shot run (slower: model loads then exits)" >&2
    fb_model="${MODEL:-sdxl}"
    OUT_NAME="${OUT:-${fb_model}-$(date +%Y%m%d-%H%M%S).png}"
    OUT_PATH="${OUT_DIR}/${OUT_NAME}"
    set -- "$PROMPT"
    [ -n "$MODEL" ] && set -- "$@" -m "$MODEL"
    [ -n "$NEG" ]   && set -- "$@" -n "$NEG"
    [ -n "$STEPS" ] && set -- "$@" -s "$STEPS"
    [ -n "$GUID" ]  && set -- "$@" -g "$GUID"
    [ -n "$SEED" ]  && set -- "$@" --seed "$SEED"
    [ -n "$W" ]     && set -- "$@" -W "$W"
    [ -n "$H" ]     && set -- "$@" -H "$H"
    set -- "$@" -o "$OUT_NAME"
    docker compose run --rm imggen generate.py "$@"
fi
