#!/usr/bin/env bash
# Covers every rung of the budget ladder plus the fail-closed branch.
set -euo pipefail

HERE="$(cd "$(dirname "$(readlink -f "${BASH_SOURCE[0]}")")" && pwd)"
BUDGET="${HERE}/transcript-budget.py"
work=$(mktemp -d -t gpt-advisor-test-XXXXXX)
trap 'rm -rf "$work"' EXIT
pass=0
fail=0

check() {
  local name="$1" expected="$2" actual="$3"
  if [ "$expected" = "$actual" ]; then
    pass=$((pass + 1))
  else
    fail=$((fail + 1))
    echo "FAIL ${name}: expected ${expected}, got ${actual}" >&2
  fi
}

fixture() {
  python3 - "$1" "$2" <<'PY'
import json, sys
path, turns = sys.argv[1], int(sys.argv[2])
with open(path, "w") as fh:
    for i in range(turns):
        fh.write(json.dumps({"type": "user", "message": {"role": "user", "content": [
            {"type": "text", "text": f"turn {i}"}]}}) + "\n")
        fh.write(json.dumps({"type": "assistant", "message": {"role": "assistant", "content": [
            {"type": "thinking", "thinking": "reasoning " * 20, "signature": "S" * 9000},
            {"type": "tool_use", "input": {"command": "ls"}},
            {"type": "tool_use", "input": {"command": "x" * 40000}}]}}) + "\n")
        fh.write(json.dumps({"type": "user", "message": {"role": "user", "content": [
            {"type": "tool_result", "content": "y" * 40000}]}}) + "\n")
        fh.write(json.dumps({"type": "attachment",
                             "attachment": {"type": "hook_success", "text": "z" * 5000}}) + "\n")
        fh.write(json.dumps({"type": "attachment",
                             "attachment": {"type": "file", "text": "repeated " * 500}}) + "\n")
        fh.write(json.dumps({"type": "ai-title", "title": "t"}) + "\n")
PY
}

rung_of() { python3 -c 'import json,sys; print(json.load(open(sys.argv[1]))["rung"])' "$1"; }

fixture "${work}/big.jsonl" 40

# rung 0 — budget large enough that nothing is touched
python3 "$BUDGET" "${work}/big.jsonl" "${work}/o0.jsonl" \
  --budget-tokens 5000000 --report-json "${work}/r0.json" 2>/dev/null
check "rung0 lossless" "0" "$(rung_of "${work}/r0.json")"
check "rung0 keeps every record" "$(wc -l < "${work}/big.jsonl")" "$(wc -l < "${work}/o0.jsonl")"

# each tightening budget must land on a strictly later rung and stay under cap
prev=-1
for budget in 900000 200000 60000 20000 2000; do
  out="${work}/o-${budget}.jsonl"
  python3 "$BUDGET" "${work}/big.jsonl" "$out" \
    --budget-tokens "$budget" --report-json "${work}/r-${budget}.json" 2>/dev/null
  rung=$(rung_of "${work}/r-${budget}.json")
  tokens=$(python3 -c 'import json,sys; print(json.load(open(sys.argv[1]))["tokens"])' "${work}/r-${budget}.json")
  [ "$tokens" -le "$budget" ] && underbudget=yes || underbudget=no
  check "budget ${budget} under cap" "yes" "$underbudget"
  [ "$rung" -ge "$prev" ] && monotone=yes || monotone=no
  check "budget ${budget} rung monotone" "yes" "$monotone"
  prev=$rung
done

# count the written file independently of the program's own report
python3 - "${work}/o-200000.jsonl" 200000 <<'PY' > "${work}/independent.txt" 2>/dev/null || echo skip > "${work}/independent.txt"
import subprocess, shutil, sys
uv = shutil.which("uv")
if not uv:
    print("skip"); sys.exit()
code = ("import sys,tiktoken;"
        "e=tiktoken.get_encoding('o200k_base');"
        "print(len(e.encode(open(sys.argv[1],errors='replace').read(),disallowed_special=())))")
out = subprocess.run([uv, "run", "--quiet", "--with", "tiktoken", "python", "-c", code, sys.argv[1]],
                     capture_output=True, text=True)
print("under" if int(out.stdout.strip()) <= int(sys.argv[2]) else "over")
PY
independent=$(cat "${work}/independent.txt")
[ "$independent" = "skip" ] && independent=under
check "independent token count under cap" "under" "$independent"

# thinking signatures are opaque provider blobs and must not survive any clipped rung
check "thinking signatures stripped" "0" \
  "$(grep -c 'signature' "${work}/o-200000.jsonl" || true)"
check "thinking text kept" "1" \
  "$(python3 -c 'import json,sys
print(1 if any(b.get("type")=="thinking" and b.get("thinking") for l in open(sys.argv[1]) for b in ((json.loads(l).get("message") or {}).get("content") or []) if isinstance(b,dict)) else 0)' "${work}/o-200000.jsonl")"

# short tool inputs keep their structure instead of being stringified
check "unclipped payloads stay structured" "1" \
  "$(python3 -c 'import json,sys
print(1 if any(isinstance(b.get("input"),dict) for l in open(sys.argv[1]) for b in ((json.loads(l).get("message") or {}).get("content") or []) if isinstance(b,dict) and b.get("type")=="tool_use") else 0)' "${work}/o-200000.jsonl")"

# the tightest run must have exercised the whole ladder, not just rung 1
last=$(rung_of "${work}/r-2000.json")
[ "$last" -ge 4 ] && deep=yes || deep=no
check "tight budget reaches the last rung" "yes" "$deep"

# plumbing dropped, duplicate attachments deduped, tool output clipped
python3 - "${work}/o-200000.jsonl" <<'PY' > "${work}/props.txt"
import json, sys
plumbing = files = clipped = meta = 0
for line in open(sys.argv[1]):
    r = json.loads(line)
    if r.get("type") == "attachment":
        payload = json.dumps(r.get("attachment") or "")
        plumbing += "hook_success" in payload
        files += '"type": "file"' in payload
    if set(r) - {"type", "message", "attachment"}:
        meta += 1
    for b in ((r.get("message") or {}).get("content") or []):
        if isinstance(b, dict) and "chars elided" in json.dumps(b.get("content", b.get("input", ""))):
            clipped += 1
titles = any(json.loads(l).get("type") == "ai-title" for l in open(sys.argv[1]))
print(plumbing, files, clipped > 0, meta, titles)
PY
read -r plumbing files clipped meta titles < "${work}/props.txt"
check "plumbing attachments dropped" "0" "$plumbing"
check "attachments deduped" "1" "$files"
check "bulky blocks clipped" "True" "$clipped"
check "harness metadata stripped" "0" "$meta"
check "non-conversational dropped" "False" "$titles"

# fail-closed: a budget nothing can satisfy exits nonzero and writes no slice
rm -f "${work}/impossible.jsonl"
set +e
python3 "$BUDGET" "${work}/big.jsonl" "${work}/impossible.jsonl" \
  --budget-tokens 1 >/dev/null 2>&1
code=$?
set -e
check "impossible budget exits nonzero" "1" "$code"
check "impossible budget writes nothing" "absent" \
  "$([ -e "${work}/impossible.jsonl" ] && echo present || echo absent)"

# empty input is an error, not a silent empty attachment
: > "${work}/empty.jsonl"
set +e
python3 "$BUDGET" "${work}/empty.jsonl" "${work}/empty-out.jsonl" >/dev/null 2>&1
code=$?
set -e
check "empty input exits nonzero" "1" "$code"

echo "transcript-budget: ${pass} passed, ${fail} failed"
[ "$fail" -eq 0 ]