#!/usr/bin/env bash
# Axiom query formatter - compact, grepable, token-efficient
# Usage: ... | axiom-query-fmt [--raw|--full]

set -euo pipefail

FULL=false
for arg in "$@"; do
  case "$arg" in
    --raw) cat; exit 0 ;;
    --full) FULL=true ;;
  esac
done

INPUT=$(cat)

if ! echo "$INPUT" | jq -e '.tables' >/dev/null 2>&1; then
  msg=$(echo "$INPUT" | jq -r '.message // empty' 2>/dev/null)
  echo "error: ${msg:-invalid response}" >&2
  exit 1
fi

echo "$INPUT" | jq -r --argjson full "$FULL" '
def fmt:
  if . == null then empty
  elif type == "boolean" then (if . then "true" else "false" end)
  elif type == "number" then
    if . == (. | floor) then tostring
    else ((. * 100 | floor) / 100 | tostring)
    end
  elif type == "string" then
    if (. | length) > 120 and ($full | not) then
      "\"" + .[0:100] + "...[+" + ((. | length) - 100 | tostring) + " chars]\""
    elif . | test("\\s") then "\"" + . + "\""
    else .
    end
  elif type == "array" then "[" + (length | tostring) + "]"
  elif type == "object" then "{" + (keys | length | tostring) + "}"
  else tostring
  end;

.tables[0] as $t |
($t.fields | map(.name)) as $f |
($t.columns // []) as $c |
(if ($c | length) > 0 then ($c[0] | length) else 0 end) as $n |

"# \(.status.rowsMatched // 0)/\(.status.rowsExamined // 0) rows, \((.status.elapsedTime // 0) / 1000 | floor)ms",
(range($n) as $i |
  [range($f | length) as $j |
    $c[$j][$i] as $v |
    if $v == null then empty
    else "\($f[$j])=\($v | fmt)"
    end
  ] | join(" ")
)
'
