#!/bin/bash
# query_gemini.sh - Query Gemini via the Antigravity CLI (`agy`) or Google AI API
#
# CLI mode uses `agy` (Antigravity CLI), the successor to the deprecated Gemini
# CLI (transitioned 2026-06-18). Models are addressed by display name and the
# CLI prints the model's response as plain text (no JSON envelope wrapper).
#
# Usage: ./query_gemini.sh "question" [context_file] [output_file]
#
# Environment variables:
#   GEMINI_MODEL - agy model display name (default: "Gemini 3.7 Flash (High)")
#   GEMINI_TIMEOUT - Timeout in seconds (default: 180)
#   GEMINI_USE_API - Use Google AI API mode instead of the agy CLI (default: false)
#   GEMINI_API_MODEL - API model ID for API mode (default: gemini-3.1-pro-preview)
#   GEMINI_API_KEY - API key for API mode
#   GEMINI_REASONING_EFFORT - Optional CLI reasoning effort
#   ENABLE_PERSONA - Enable "The Architect" persona (default: true)

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
source "$SCRIPT_DIR/lib/common.sh"
source "$SCRIPT_DIR/lib/personas.sh"

# --- Parameters ---
QUERY="${1:-}"
CONTEXT_FILE="${2:-}"
OUTPUT_FILE="${3:-/tmp/gemini_response.json}"

# --- Configuration ---
ENABLE_PERSONA="${ENABLE_PERSONA:-true}"
CONSULTANT_NAME="Gemini"

# --- Build query ---
FULL_QUERY=$(build_full_query "$QUERY" "$CONTEXT_FILE")
validate_query "$FULL_QUERY" "Gemini" || exit 1

# --- Add persona if enabled ---
if [[ "$ENABLE_PERSONA" == "true" ]]; then
    FULL_QUERY=$(build_query_with_persona "$CONSULTANT_NAME" "$FULL_QUERY")
fi

# --- Timestamp for metadata ---
START_TIME=$(get_timestamp_ms)

# --- Execution (CLI or API mode) ---
TEMP_OUTPUT=$(mktemp)
trap 'rm -f "$TEMP_OUTPUT" "${TEMP_OUTPUT}.err"' EXIT
MODEL_IDENTITY_SOURCE="requested-only"
EFFECTIVE_MODEL=""
GEMINI_PROBE_REASON=""

gemini_cli_exposes_requested_model() {
    local inventory requested_norm
    if ! inventory=$(run_with_timeout 10 "${AGY_ENV_PREFIX[@]}" "$GEMINI_CMD" models 2>&1); then
        GEMINI_PROBE_REASON="Antigravity model inventory timed out, failed, or requires login"
        printf '%s\n' "$inventory" >&2
        return 1
    fi
    requested_norm=$(printf '%s' "$GEMINI_MODEL" | tr '[:upper:]' '[:lower:]' | tr -cd '[:alnum:]')
    if ! printf '%s\n' "$inventory" | awk -v requested="$requested_norm" '
        {
            line = tolower($0)
            gsub(/[^[:alnum:]]/, "", line)
            if (index(line, requested) > 0) found = 1
        }
        END { exit(found ? 0 : 1) }
    '; then
        GEMINI_PROBE_REASON="Antigravity model inventory does not include $GEMINI_MODEL"
        return 1
    fi
}

if is_api_mode "gemini"; then
    # --- API Mode ---
    log_api_mode_status "gemini"
    validate_api_mode "gemini" || exit 1

    source "$SCRIPT_DIR/lib/api_query.sh"

    # API mode addresses the Google AI endpoint, which needs an API model ID
    # (not an agy display name like "Gemini 3.1 Pro (High)").
    # Keep the failure inside an explicit conditional: a bare call whose
    # function returns non-zero aborts the script under `set -e` before
    # exit_code is read, so no error-response envelope is ever written and
    # the output file is left empty. Same guard run_query uses.
    if run_api_mode_query \
            "$CONSULTANT_NAME" \
            "$GEMINI_API_MODEL" \
            "$FULL_QUERY" \
            "$TEMP_OUTPUT" \
            "$GEMINI_TIMEOUT_SECONDS"; then
        exit_code=0
    else
        exit_code=$?
    fi
else
    # --- CLI Mode (Antigravity CLI: agy) ---
    log_api_mode_status "gemini"
    check_command "$GEMINI_CMD" "Antigravity CLI" "curl -fsSL https://antigravity.google/cli/install.sh | bash" || exit 1
    if ! gemini_cli_exposes_requested_model; then
        log_error "[$CONSULTANT_NAME] $GEMINI_PROBE_REASON"
        exit 1
    fi
    MODEL_IDENTITY_SOURCE="capability-probed"

    # agy prints the model's response as plain text -- there is no CLI envelope
    # to unwrap. The persona instruction forces the model to emit our JSON schema
    # (some models wrap it in a ```json markdown fence; process_consultant_response
    # strips it centrally).
    # NOTE: agy's -p/--print/--prompt takes the prompt as its ARGUMENT value -- it
    # does NOT read stdin, and "-" is not a stdin sentinel. A prior `-p -` shipped
    # silently broken: agy answered a literal "-" with a generic greeting (exit 0,
    # so no error surfaced -> fallback envelope). The prompt therefore rides as the
    # -p argument. agy has --model (no -m alias) and no read-from-file flag, so a
    # very large FULL_QUERY goes through argv (ARG_MAX-bounded; fine for normal
    # contexts). stdin is /dev/null (agy ignores it; run_query's cat needs an EOF).
    # Strip SSH_* markers so agy consults the macOS Keychain (see agy_env /
    # AGY_ENV_PREFIX in lib/common.sh). Use the exec-safe prefix array — not the
    # agy_env function — because run_query dispatches through GNU timeout.
    #
    # Optional CLI effort: when GEMINI_REASONING_EFFORT is set, pass --effort.
    # agy accepts only a subset (e.g. low|medium|high); unsupported values fail
    # at the provider and surface as diagnosed errors — no local allowlist.
    GEMINI_ARGS=("${AGY_ENV_PREFIX[@]}" "$GEMINI_CMD" -p "$FULL_QUERY" --model "$GEMINI_MODEL")
    effort_ok=true
    if [[ -n "${GEMINI_REASONING_EFFORT:-}" ]]; then
        source "$SCRIPT_DIR/lib/api.sh"
        if ! cli_effort=$(validate_reasoning_effort "$GEMINI_REASONING_EFFORT" "$CONSULTANT_NAME"); then
            effort_ok=false
            exit_code=1
        else
            GEMINI_ARGS+=(--effort "$cli_effort")
        fi
    fi

    if [[ "$effort_ok" == "true" ]]; then
        # Keep failures inside an explicit conditional so `set -e` cannot skip
        # exit_code assignment (same guard as the API branch).
        if run_query \
                "Gemini" \
                "$TEMP_OUTPUT" \
                "$GEMINI_TIMEOUT_SECONDS" \
                "${GEMINI_ARGS[@]}" </dev/null; then
            exit_code=0
        else
            exit_code=$?
        fi
    fi
fi

# --- Calculate latency ---
END_TIME=$(get_timestamp_ms)
LATENCY_MS=$((END_TIME - START_TIME))

# --- Configuration for response building ---
# Record the identifier actually sent to the provider so cost lookup uses the
# matching rate table entry. CLI and API mode intentionally use different model
# namespaces.
if is_api_mode "gemini"; then
    MODEL_USED="$GEMINI_API_MODEL"
    EFFECTIVE_MODEL="${_API_RESPONSE_MODEL:-$MODEL_USED}"
    MODEL_IDENTITY_SOURCE="${_API_MODEL_IDENTITY_SOURCE:-requested-only}"
else
    MODEL_USED="$GEMINI_MODEL"
    EFFECTIVE_MODEL="$MODEL_USED"
fi
PERSONA_NAME=$(get_persona_name "$CONSULTANT_NAME")

# --- Post-processing: use shared helper ---
# No native_json_field: agy prints the model's JSON directly (top-level
# .response, possibly inside a ```json fence that process_consultant_response
# strips), so extracting ".response" here would strip a level. The old Gemini
# CLI wrapped output in {"response": "..."} and needed that argument.
if process_consultant_response "$CONSULTANT_NAME" "$MODEL_USED" "$PERSONA_NAME" \
        "$TEMP_OUTPUT" "$OUTPUT_FILE" "$exit_code" "$LATENCY_MS" "" "$FULL_QUERY" \
        "$MODEL_USED" "$MODEL_IDENTITY_SOURCE" "$EFFECTIVE_MODEL"; then
    :
else
    response_rc=$?
    [[ $exit_code -ne 0 ]] || exit_code=$response_rc
fi

cat "$OUTPUT_FILE"
exit $exit_code