#!/usr/bin/env bash
# ═══════════════════════════════════════════════════════════════════
# test-antigravity.sh — End-to-end Antigravity proxy test
#
# Tests: token validity → direct REST probe → proxy adapter
#
# Usage: bash ~/.local/bin/test-antigravity.sh [--verbose]
# Exit:  0 = all pass, 1 = some fail
# ═══════════════════════════════════════════════════════════════════
set -euo pipefail

VERBOSE=0
for arg in "$@"; do case "$arg" in --verbose|-v) VERBOSE=1 ;; esac; done

RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'; NC='\033[0m'
PASS=0; FAIL=0; SKIP=0; RESULTS=()
log_pass() { echo -e "  ${GREEN}PASS${NC} $1"; ((PASS++)); RESULTS+=("PASS $1"); }
log_fail() { echo -e "  ${RED}FAIL${NC} $1"; ((FAIL++)); RESULTS+=("FAIL $1"); }
log_skip() { echo -e "  ${YELLOW}SKIP${NC} $1"; ((SKIP++)); RESULTS+=("SKIP $1"); }

TOKEN_PATH="$HOME/.cache/codex-proxy/google-antigravity-oauth-token.json"
[ ! -f "$TOKEN_PATH" ] && { echo "ERROR: No token file. Login via GUI first."; exit 1; }

ACCESS_TOKEN=$(python3 -c "
import json, os, sys, time, urllib.request, urllib.parse
tp = os.path.expanduser('~/.cache/codex-proxy/google-antigravity-oauth-token.json')
d = json.load(open(tp))
if d.get('expires_at', 0) > time.time(): print(d['access_token']); sys.exit(0)
cid, cs, rt = d.get('client_id',''), d.get('client_secret',''), d.get('refresh_token','')
if not all([cid, cs, rt]): print('ERROR'); sys.exit(1)
data = urllib.parse.urlencode({'client_id':cid,'client_secret':cs,'refresh_token':rt,'grant_type':'refresh_token'}).encode()
resp = urllib.request.urlopen(urllib.request.Request('https://oauth2.googleapis.com/token', data=data), timeout=15)
tok = json.loads(resp.read()); d.update(tok); d['expires_at'] = time.time() + tok.get('expires_in',3600)
json.dump(d, open(tp,'w')); print(tok.get('access_token','ERROR'))
" 2>&1) || true
[[ "$ACCESS_TOKEN" == ERROR* ]] || [ -z "$ACCESS_TOKEN" ] && { echo "ERROR: Token refresh failed: $ACCESS_TOKEN"; exit 1; }

PROJECT_ID=$(python3 -c "import json; print(json.load(open('$TOKEN_PATH')).get('project_id',''))")
[ -z "$PROJECT_ID" ] && { echo "ERROR: No project_id"; exit 1; }

echo "═══════════════════════════════════════════════════════════════"
echo " Antigravity E2E Test Suite"
echo "═══════════════════════════════════════════════════════════════"
echo " Project: $PROJECT_ID  Token: ${ACCESS_TOKEN:0:20}..."

# ── Test 1: Token validity ────────────────────────────────────────
echo ""; echo "─── Test 1: Token Validity ───"
HTTP=$(curl -s -o /dev/null -w "%{http_code}" -H "Authorization: Bearer $ACCESS_TOKEN" \
    "https://www.googleapis.com/oauth2/v1/userinfo" --max-time 5)
[ "$HTTP" = "200" ] && log_pass "Token valid" || log_fail "Token invalid (HTTP $HTTP)"

# ── Test 2: Direct REST probe (prod first, fast timeout) ─────────
echo ""; echo "─── Test 2: Direct REST Endpoint Probe ───"
ENDPOINTS=(
    "https://cloudcode-pa.googleapis.com"
    "https://daily-cloudcode-pa.sandbox.googleapis.com"
    "https://autopush-cloudcode-pa.sandbox.googleapis.com"
)
MODELS=("gemini-3-flash")
BEST_EP=""; BEST_MODEL=""

for model in "${MODELS[@]}"; do
    for ep in "${ENDPOINTS[@]}"; do
        ep_s=$(echo "$ep" | sed 's|https://||;s|.googleapis.com||')
        RESP=$(curl -s -w "\n%{http_code}" -X POST "${ep}/v1internal:generateContent" \
            -H "Content-Type: application/json" \
            -H "Authorization: Bearer $ACCESS_TOKEN" \
            -H "User-Agent: Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Antigravity/2.0.6 Chrome/138.0.7204.235 Electron/37.3.1 Safari/537.36" \
            -H 'Client-Metadata: {"ideType":"ANTIGRAVITY","platform":"LINUX","pluginType":"GEMINI"}' \
            -d "{\"project\":\"$PROJECT_ID\",\"model\":\"$model\",\"requestType\":\"agent\",\"userAgent\":\"antigravity/2.0.6 linux/x64\",\"requestId\":\"t$(date +%s)\",\"request\":{\"contents\":[{\"role\":\"user\",\"parts\":[{\"text\":\"Say hi\"}]}],\"sessionId\":\"t$(date +%s%N)\",\"generationConfig\":{\"maxOutputTokens\":256}}}" \
            --connect-timeout 5 --max-time 20 2>&1)
        HTTP=$(echo "$RESP" | tail -1); BODY=$(echo "$RESP" | sed '$d')
        if [ "$HTTP" = "200" ]; then
            TEXT=$(echo "$BODY" | python3 -c "
import sys, json
try:
    d = json.load(sys.stdin)
    parts = d.get('response',{}).get('candidates',[{}])[0].get('content',{}).get('parts',[])
    texts = [p['text'] for p in parts if 'text' in p and p['text']]
    print(' '.join(texts)[:80] if texts else 'EMPTY')
except: print('EMPTY')" 2>/dev/null)
            if [ "$TEXT" != "EMPTY" ] && ! echo "$TEXT" | grep -qi "no longer supported"; then
                log_pass "$model @ ${ep_s} → \"$TEXT\""
                [ -z "$BEST_EP" ] && BEST_EP="$ep" && BEST_MODEL="$model"
            else
                log_fail "$model @ ${ep_s} → 200 but empty/deprecated"
            fi
        else
            ERR=$(echo "$BODY" | python3 -c "
import sys, json
try: print(json.load(sys.stdin).get('error',{}).get('status','')[:50])
except: pass" 2>/dev/null)
            log_skip "$model @ ${ep_s} → $HTTP $ERR"
        fi
    done
done

# ── Test 3: Proxy adapter (start proxy, test /responses) ──────────
echo ""; echo "─── Test 3: Proxy Adapter (end-to-end) ───"

TEST_PORT=$(python3 -c "import socket; s=socket.socket(); s.bind(('',0)); print(s.getsockname()[1]); s.close()")
PROXY_API_KEY="test-$RANDOM"

find /home/roman/.local/bin -name "__pycache__" -type d -exec rm -rf {} + 2>/dev/null || true

PROXY_PID=""
PROXY_PORT=$TEST_PORT PROXY_API_KEY=$PROXY_API_KEY PROXY_BACKEND=gemini-oauth-antigravity \
    PROXY_TARGET_URL=https://cloudcode-pa.googleapis.com \
    python3 /home/roman/.local/bin/translate-proxy.py >/tmp/antigravity-test-proxy.log 2>&1 &
PROXY_PID=$!

cleanup() { kill $PROXY_PID 2>/dev/null || true; wait $PROXY_PID 2>/dev/null || true; }
trap cleanup EXIT

sleep 3
if ! kill -0 $PROXY_PID 2>/dev/null; then
    log_fail "Proxy failed to start (port $TEST_PORT)"
    cat /tmp/antigravity-test-proxy.log 2>/dev/null | tail -5
else
    log_pass "Proxy started on :$TEST_PORT"

    # /v1/models
    HTTP=$(curl -s -o /dev/null -w "%{http_code}" -H "Authorization: Bearer $PROXY_API_KEY" \
        "http://127.0.0.1:$TEST_PORT/v1/models" --max-time 5)
    [ "$HTTP" = "200" ] && log_pass "/v1/models → 200" || log_fail "/v1/models → $HTTP"

    # /responses (non-stream)
    RESP_HTTP=$(curl -s -w "%{http_code}" -o /tmp/antigravity-test-response.json \
        -X POST "http://127.0.0.1:$TEST_PORT/responses" \
        -H "Content-Type: application/json" \
        -H "Authorization: Bearer $PROXY_API_KEY" \
        -d '{
            "model":"gemini-3.5-flash-high",
            "stream":false,
            "input":[{"type":"message","role":"user","content":[{"type":"input_text","text":"Say hello in exactly 3 words"}]}],
            "tools":[{"type":"function","name":"test_tool","description":"test","parameters":{"type":"object","properties":{"cmd":{"type":"string"}}}}],
            "instructions":"You are a helpful assistant.",
            "max_output_tokens":256
        }' --connect-timeout 10 --max-time 60 2>&1)

    if [ "$RESP_HTTP" = "200" ]; then
        TEXT=$(python3 -c "
import json
d = json.load(open('/tmp/antigravity-test-response.json'))
out = d.get('output', [])
texts = []
for item in out:
    for p in (item.get('content', []) if isinstance(item, dict) else []):
        if isinstance(p, dict): texts.append(p.get('text', ''))
print(' '.join(t for t in texts if t).strip()[:120] or 'EMPTY')
" 2>/dev/null)
        if [ "$TEXT" = "EMPTY" ]; then
            log_fail "Proxy /responses → 200 but EMPTY"
        else
            log_pass "Proxy /responses → 200: \"$TEXT\""
        fi
    else
        ERR=$(python3 -c "
import json; d = json.load(open('/tmp/antigravity-test-response.json'))
print(d.get('error',{}).get('message','')[:120])" 2>/dev/null || echo "unknown")
        log_fail "Proxy /responses → $RESP_HTTP: $ERR"
    fi

    # Verify model resolution in logs
    if grep -q "model resolved: gemini-3.5-flash-high -> gemini-3-flash" /tmp/antigravity-test-proxy.log; then
        log_pass "Model resolution: gemini-3.5-flash-high → gemini-3-flash"
    else
        log_fail "Model resolution not found in proxy logs"
    fi

    [ "$VERBOSE" = "1" ] && cat /tmp/antigravity-test-proxy.log
fi

# ── Summary ───────────────────────────────────────────────────────
echo ""
echo "═══════════════════════════════════════════════════════════════"
echo " Results: $PASS passed, $FAIL failed, $SKIP skipped"
echo "═══════════════════════════════════════════════════════════════"
[ -n "$BEST_EP" ] && echo -e " ${GREEN}Best direct:${NC} $BEST_MODEL @ $BEST_EP"

if [ "$FAIL" -gt 0 ]; then
    echo -e "\n${RED}FAILED — Do NOT push until all tests pass${NC}"
    for r in "${RESULTS[@]}"; do echo "$r" | grep -q "^FAIL" && echo "  $r"; done
    exit 1
else
    echo -e "\n${GREEN}ALL TESTS PASSED — Safe to push${NC}"
    exit 0
fi