#!/usr/bin/env bash
set -euo pipefail
# POS: ai gemini — Chat with Google Gemini (ask, chat, models, sessions)
# POS_SUBCMDS: ask chat models sessions
# POS_FLAGS: --model --session --system
# POS_CONFIG: ai | ai.env | AI_GEMINI_API_KEY=secret:API key from aistudio.google.com | AI_GEMINI_MODEL=:Model id (default gemini-2.5-flash)

source "$(dirname "$0")/../lib/common.sh" 2>/dev/null || source "$(dirname "$0")/common.sh"

CONFIG_FILE="$CONFIG_DIR/ai.env"
API="https://generativelanguage.googleapis.com/v1beta"
DEFAULT_MODEL="gemini-2.5-flash"
SESSION_DIR="$HOME/.local/share/linux_post_install/ai"
SESSION=""
SYSTEM_PROMPT=""
MAX_SESSION_TURNS=40

usage() {
    cat <<EOF
Usage: pos ai gemini <subcommand> [--model <id>] [--session <name>] [--system <text>]

Chat with Google Gemini via the REST API (generativelanguage.googleapis.com).

Subcommands:
  ask "<prompt>"    One-shot answer; prints ONLY the answer text to stdout
                    (pipe/script/Telegram-friendly). The prompt may also be
                    piped in via stdin when no argument is given. With
                    --session, previous turns are sent as context.
  chat              Interactive multi-turn conversation.
  models            List models that support generateContent.
  sessions          List persistent sessions / clear one:
                    'sessions' and 'sessions reset <name>'.

Options:
  --model <id>      Override the model for this invocation.
  --session <name>  Persistent memory: ask/chat remember prior turns in
                    ~/.local/share/linux_post_install/ai/<name>.json
                    (capped at $MAX_SESSION_TURNS turns). ask without
                    --session stays one-shot.
  --system <text>   System instruction sent with every turn (kept out of the
                    session file), e.g. "Reply like a friendly Telegram chat".
  -h|--help         This help.

Config:  $CONFIG_FILE  (edit with 'pos config ai')
  AI_GEMINI_API_KEY   API key from aistudio.google.com (required)
  AI_GEMINI_MODEL     Model id (default $DEFAULT_MODEL)

Examples:
  pos ai gemini ask "Explain DNS in one line"
  echo "summarize this log" | pos ai gemini ask
  pos ai gemini chat
  pos ai gemini models
  pos ai gemini ask --model gemini-2.5-flash "hi"
  pos ai gemini ask --session work "my name is joe"
  pos ai gemini ask --session work "what is my name?"   # remembers
  pos ai gemini sessions
  pos ai gemini sessions reset work
EOF
    exit 0
}

# ── ai.env loader ──────────────────────────────────────────────
load_config() { load_tool_config "$CONFIG_FILE"; }

require_key() {
    load_config
    [ -n "${AI_GEMINI_API_KEY:-}" ] || err "No Gemini API key — run 'pos config ai'"
}

resolve_model() {
    if [ -n "${MODEL_OVERRIDE:-}" ]; then
        printf '%s' "$MODEL_OVERRIDE"
    elif [ -n "${AI_GEMINI_MODEL:-}" ]; then
        printf '%s' "$AI_GEMINI_MODEL"
    else
        printf '%s' "$DEFAULT_MODEL"
    fi
}

# ── Persistent session memory ───────────────────────────────────
# History lives as a Gemini "contents" JSON document per session name under
# SESSION_DIR. Names are sanitized to [A-Za-z0-9_-]; ask/chat only touch the
# session layer when --session is given (otherwise they stay stateless).
session_file() {
    local name="${1:-$SESSION}"
    name="${name//[^A-Za-z0-9_-]/_}"
    printf '%s/%s.json' "$SESSION_DIR" "$name"
}

session_load() {
    [ -n "$SESSION" ] || { printf '{"contents":[]}'; return 0; }
    local f
    f="$(session_file)"
    if [ -s "$f" ] && jq -e '.contents' "$f" >/dev/null 2>&1; then
        cat "$f"
    else
        printf '{"contents":[]}'
    fi
}

session_save() {
    [ -n "$SESSION" ] || return 0
    local f tmp
    f="$(session_file)"
    mkdir -p "$SESSION_DIR"
    tmp="$(mktemp)"
    printf '%s\n' "$1" >"$tmp"
    mv "$tmp" "$f"
    chmod 600 "$f"
}

# Append a turn and prune to the last MAX_SESSION_TURNS entries. stdout = JSON.
session_push() {
    local contents="$1" role="$2" text="$3"
    printf '%s' "$contents" | jq -c --arg r "$role" --arg t "$text" \
        '.contents += [{role:$r, parts:[{text:$t}]}] | .contents |= .[-'"$MAX_SESSION_TURNS"':]'
}

# One generateContent call. $1 = model, $2 = contents JSON, $3 = optional
# system instruction (added as systemInstruction, not stored in the session).
# stdout = the answer text on success; an error message on failure (exit 1).
gemini_generate() {
    local model="$1" contents="$2" system="${3:-}" body
    body="$contents"
    if [ -n "$system" ]; then
        body="$(printf '%s' "$contents" | jq -c --arg s "$system" \
            '. + {systemInstruction:{role:"system",parts:[{text:$s}]}}')"
    fi
    local resp code body_out errmsg
    resp="$(curl -sS -m 60 -X POST "${API}/models/${model}:generateContent" \
        -H "x-goog-api-key: ${AI_GEMINI_API_KEY}" \
        -H "Content-Type: application/json" \
        --write-out $'\n%{http_code}' \
        --data "$body")" || { echo "request failed (curl exit $?)" >&2; return 1; }
    code="${resp##*$'\n'}"
    body_out="${resp%$'\n'*}"
    if [ "$code" != "200" ]; then
        errmsg="$(printf '%s' "$body_out" | jq -r '.error.message // empty' 2>/dev/null || true)"
        echo "API error $code${errmsg:+: $errmsg}" >&2
        return 1
    fi
    printf '%s' "$body_out" | jq -r '[.candidates[0].content.parts[]?.text] | join("")'
}

cmd_ask() {
    local prompt="" contents out
    if [ $# -gt 0 ]; then
        prompt="$*"
    elif [ ! -t 0 ]; then
        prompt="$(cat)"
    fi
    [ -n "$prompt" ] || err "No prompt given — usage: pos ai gemini ask \"<prompt>\""
    require_key
    if [ -n "$SESSION" ]; then
        contents="$(session_load)"
        contents="$(session_push "$contents" user "$prompt")"
    else
        contents="$(jq -nc --arg t "$prompt" '{contents:[{role:"user",parts:[{text:$t}]}]}')"
    fi
    if ! out="$(gemini_generate "$(resolve_model)" "$contents" "$SYSTEM_PROMPT" 2>&1)"; then
        err "$out"
    fi
    if [ -n "$SESSION" ]; then
        contents="$(session_push "$contents" model "$out")"
        session_save "$contents"
    fi
    printf '%s\n' "$out"
}

cmd_chat() {
    [ $# -eq 0 ] || err "Unexpected argument for chat: $*"
    local model contents text answer
    model="$(resolve_model)"
    require_key
    if [ -n "$SESSION" ]; then
        contents="$(session_load)"
        printf 'session: %s (resumed %s prior turns)\n' "$SESSION" "$(printf '%s' "$contents" | jq -r '.contents | length')"
    else
        contents='{"contents":[]}'
    fi
    trap 'echo; echo "bye"; exit 0' INT
    echo "Gemini · ${model} — type a message; q=quit, /reset=clear history"
    while true; do
        printf '> '
        IFS= read -r text || break
        case "$text" in
            "" ) continue ;;
            q|Q|quit|exit) echo; echo "bye"; return 0 ;;
            /reset)
                contents='{"contents":[]}'
                [ -n "$SESSION" ] && session_save "$contents"
                echo "[history cleared]"
                continue ;;
        esac
        contents="$(session_push "$contents" user "$text")"
        if ! answer="$(gemini_generate "$model" "$contents" "$SYSTEM_PROMPT" 2>&1)"; then
            warn "AI error: $answer"
            continue
        fi
        contents="$(session_push "$contents" model "$answer")"
        [ -n "$SESSION" ] && session_save "$contents"
        printf '\n%s\n\n' "$answer"
    done
    echo
    return 0
}

cmd_sessions() {
    local action="${1:-list}" name f n
    case "$action" in
        list|"")
            [ -d "$SESSION_DIR" ] || { echo "no sessions"; return 0; }
            local found=0
            for f in "$SESSION_DIR"/*.json; do
                [ -f "$f" ] || continue
                found=1
                n="$(jq -r '.contents | length' "$f" 2>/dev/null || echo 0)"
                printf '  %-32s %s turns\n' "$(basename "$f" .json)" "${n:-0}"
            done
            [ "$found" -eq 1 ] || echo "no sessions"
            ;;
        reset)
            [ $# -ge 2 ] || err "usage: pos ai gemini sessions reset <name>"
            name="$2"
            if rm -f "$(session_file "$name")"; then
                ok "session '$name' cleared"
            fi
            ;;
        *) err "Unknown sessions subcommand '$action' (list | reset <name>)" ;;
    esac
}

cmd_models() {
    [ $# -eq 0 ] || err "Unexpected argument for models: $*"
    local model resp code body m
    model="$(resolve_model)"
    require_key
    resp="$(curl -sS -m 30 -G "${API}/models" \
        -H "x-goog-api-key: ${AI_GEMINI_API_KEY}" \
        --data-urlencode "pageSize=1000" \
        --write-out $'\n%{http_code}')" || err "request failed (curl exit $?)"
    code="${resp##*$'\n'}"
    body="${resp%$'\n'*}"
    if [ "$code" != "200" ]; then
        err "API error $code: $(printf '%s' "$body" | jq -r '.error.message // empty')"
    fi
    local list
    list="$(printf '%s' "$body" | jq -r '.models[]? | select((.supportedGenerationMethods // []) | index("generateContent")) | .name' | sed 's#^models/##' | sort)"
    echo "Gemini models (generateContent-capable):"
    while IFS= read -r m; do
        [ -n "$m" ] || continue
        if [ "$m" = "$model" ]; then
            printf '  %-32s <- default\n' "$m"
        else
            printf '  %-32s\n' "$m"
        fi
    done <<< "$list"
    if ! grep -qxF "$model" <<< "$list"; then
        warn "configured default '$model' is not in the list — set AI_GEMINI_MODEL"
    fi
}

# ── Parse flags + subcommand ────────────────────────────────────
MODEL_OVERRIDE=""
cmd=""
args=()
while [ $# -gt 0 ]; do
    case "$1" in
        -h|--help) usage ;;
        --model)
            [ $# -ge 2 ] || err "--model needs a value"
            MODEL_OVERRIDE="$2"; shift 2 ;;
        --session)
            [ $# -ge 2 ] || err "--session needs a value"
            SESSION="$2"; shift 2 ;;
        --system)
            [ $# -ge 2 ] || err "--system needs a value"
            SYSTEM_PROMPT="$2"; shift 2 ;;
        -*) err "Unknown option '$1' (see --help)" ;;
        *)
            if [ -z "$cmd" ]; then
                cmd="$1"
            else
                args+=("$1")
            fi
            shift ;;
    esac
done

case "${cmd:-}" in
    "")     usage ;;
    ask)    cmd_ask "${args[@]}" ;;
    chat)   cmd_chat "${args[@]}" ;;
    models) cmd_models "${args[@]}" ;;
    sessions) cmd_sessions "${args[@]}" ;;
    *)      err "Unknown ai gemini subcommand '$cmd' (see --help)" ;;
esac