#!/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=""; H=""
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
    # Interactive menu only when no -m and a TTY is present.
    if [ -z "$MODEL" ] && [ -t 0 ] && [ -t 1 ]; then
        echo "which model? 1) sdxl  2) realvis  3) pony" >&2
        read -rp "choice [1]: " choice || choice=1
        case "$choice" in 2) MODEL=realvis ;; 3) MODEL=pony ;; *) MODEL=sdxl ;; esac
    fi
    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
        curl -sf -X POST "${BASE}/model" -H 'Content-Type: application/json' \
            -d "{\"name\":\"$MODEL\"}" >/dev/null
    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}"

    if curl -sf -X POST "${BASE}/generate" -H 'Content-Type: application/json' \
        -d "$req" -o "$OUT_PATH"; then
        echo "Saved ${OUT_PATH}"
    else
        echo "generate failed (see daemon logs: docker compose logs)" >&2
        rm -f "$OUT_PATH"
        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
