#!/usr/bin/env zsh
set -euo pipefail

script_dir="${0:A:h}"
source "${script_dir}/lib/common.zsh"
source "${script_dir}/lib/json.zsh"

usage='Usage: query_groq [--json] "prompt" [model]'
inferencer_parse_query_args "$usage" "$@" || exit $?
inferencer_require_command jq
inferencer_require_command curl

api_key="${GROQ_API_KEY:-}"
url="${GROQ_API_URL:-https://api.groq.com/openai/v1/chat/completions}"
model="${INFERENCER_MODEL:-${GROQ_MODEL:-groq/compound}}"
temperature="${GROQ_TEMPERATURE:-1.0}"
top_p="${GROQ_TOP_P:-1.0}"
max_tokens="${GROQ_MAX_TOKENS:-1024}"
stream="${GROQ_STREAM:-false}"

inferencer_require_api_key "$api_key" "GROQ_API_KEY"
inferencer_require_bool "$stream" "GROQ_STREAM"

body="$(
  jq -n \
    --arg prompt "$INFERENCER_PROMPT" \
    --arg model "$model" \
    --argjson temperature "$temperature" \
    --argjson top_p "$top_p" \
    --argjson max_tokens "$max_tokens" \
    --argjson stream "$stream" \
    '{
      messages: [{ role: "user", content: $prompt }],
      model: $model,
      temperature: $temperature,
      top_p: $top_p,
      stream: $stream,
      max_tokens: $max_tokens
    }'
)"

response="$(
  curl --silent --show-error --fail \
    -H "Content-Type: application/json" \
    -H "Authorization: Bearer $api_key" \
    --data "$body" \
    "$url"
)"

if inferencer_response_has_error "$response"; then
  inferencer_print_error_response "$response"
  exit 1
fi

if [[ "$INFERENCER_JSON" == true ]]; then
  printf '%s\n' "$response" | jq .
  exit 0
fi

text="$(inferencer_extract_openai_text "$response")"
[[ -n "$text" ]] || {
  print -r -- "Error: Groq API returned no message content." >&2
  inferencer_print_error_response "$response"
  exit 1
}

printf '%s\n' "$text"
