#!/bin/bash
# spend-monitor-collect.sh — Collect LLM spend data for daily monitor
# Runs as script for cron job, output fed to agent prompt for analysis
# Tries SSH+Postgres first, falls back to LiteLLM HTTPS API on failure.

echo "=== LLM Spend Report — $(date +'%b %d %Y %H:%M') ==="

SSH_OK=true
ssh -i /root/.ssh/itpp-infra -o ConnectTimeout=5 -o BatchMode=yes root@178.156.167.181 "echo ok" 2>/dev/null || SSH_OK=false

# Helper: extract admin-ai key for HTTPS fallback
get_admin_key() {
  grep -A2 'admin-ai:' ~/.hermes/config.yaml | grep api_key | awk '{print $2}'
}

if [ "$SSH_OK" = true ]; then
  SSH_BASE="ssh -i /root/.ssh/itpp-infra -o ConnectTimeout=5 -o BatchMode=yes root@178.156.167.181"
  PSQL="docker exec -e PGPASSWORD=litellm litellm_postgres psql -U litellm -d litellm_db -t -c"
  QUERY="$SSH_BASE $PSQL"

  # 1. LiteLLM key spend (cumulative by key)
  echo ""
  echo "--- LiteLLM Key Spend (cumulative) ---"
  $QUERY "SELECT key_name, ROUND(spend::numeric,2) as total FROM \"LiteLLM_VerificationToken\" WHERE spend > 0 ORDER BY spend DESC;"

  # 2. Last 24h total spend
  echo ""
  echo "--- Last 24h Total Spend ---"
  $QUERY "SELECT CAST(SUM(spend) AS numeric(10,2)) as last_24h FROM \"LiteLLM_SpendLogs\" WHERE \"startTime\" > NOW() - INTERVAL '24 hours';"

  # 3. DeepSeek models last 24h
  echo ""
  echo "--- DeepSeek Models Last 24h ---"
  $QUERY "SELECT model, ROUND(SUM(spend)::numeric,4) as cost FROM \"LiteLLM_SpendLogs\" WHERE \"startTime\" > NOW() - INTERVAL '24 hours' AND model LIKE '%deepseek%' GROUP BY model ORDER BY cost DESC;"

  # 4. Top models last 24h
  echo ""
  echo "--- All Models Last 24h (top 10) ---"
  $QUERY "SELECT model, ROUND(SUM(spend)::numeric,4) as cost FROM \"LiteLLM_SpendLogs\" WHERE \"startTime\" > NOW() - INTERVAL '24 hours' GROUP BY model ORDER BY cost DESC LIMIT 10;"

else
  echo "[SSH UNAVAILABLE — using HTTPS fallback]"

  KEY=$(grep -A2 'admin-ai:' ~/.hermes/config.yaml | grep api_key | awk '{print $2}')
  [ -z "$KEY" ] && echo "ERROR: No admin-ai API key found in config" && exit 1

  # 1. Per-key cumulative spend via HTTPS
  echo ""
  echo "--- LiteLLM Key Spend (cumulative, via HTTPS) ---"
  curl -s --max-time 10 "https://admin-ai.itpropartner.com/spend/keys" \
    -H "Authorization: Bearer $KEY" | python3 -c "
import sys, json
data = json.load(sys.stdin)
for k in data:
    name = k.get('key_name', 'unknown')
    alias = k.get('key_alias', '')
    spend = k.get('spend', 0)
    label = alias if alias else name
    print(f' {label:30s} | {spend:8.2f}')
"

  # 2. Global total
  echo ""
  echo "--- Global Total Spend (via HTTPS) ---"
  curl -s --max-time 10 "https://admin-ai.itpropartner.com/global/spend" \
    -H "Authorization: Bearer $KEY" | python3 -c "
import sys, json
d = json.load(sys.stdin)
print(f' Total cumulative: \${d.get(\"spend\",0):.2f}')
print(f' Max budget:      {\"Unlimited\" if d.get(\"max_budget\",0) == 0.0 else f\"\${d[\"max_budget\"]:.2f}\"}')
"

  # 3. Per-tag breakdown
  echo ""
  echo "--- Spend by Credential/Tag (via HTTPS, cumulative) ---"
  curl -s --max-time 10 "https://admin-ai.itpropartner.com/spend/tags" \
    -H "Authorization: Bearer $KEY" | python3 -c "
import sys, json
data = json.load(sys.stdin)
for t in sorted(data, key=lambda x: x['total_spend'], reverse=True):
    tag = t.get('individual_request_tag', '')
    spend = t.get('total_spend', 0)
    count = t.get('log_count', 0)
    print(f' {tag:45s} | \${spend:>8.2f} | {count:>6} reqs')
"

  echo ""
  echo "[NOTE: HTTPS API provides cumulative (lifetime) data only. 24h model-level breakdown requires SSH+Postgres access.]"
fi