#!/usr/bin/env bash
# Kaitaku diagnostic bundler.
#
# Usage:
#   curl -sSL https://kaitaku.ai/report.sh | bash -s -- <token> \
#       [--image NAME:TAG] [--container NAME|ID] [--log /path/to/file.log]
#
# Auto-detects the runtime environment:
#   - bare-metal + Docker        → docker inspect/logs (filtered by --image)
#   - inside a container         → /proc-based process probe + known log paths
#   - bare-metal w/o Docker      → journalctl + common log paths
#
# Collected sources (best-effort, every block in `|| true`):
#   manifest.txt          summary: env, what was collected/skipped, args
#   sysinfo.txt           uname, OS release, free, df, filtered env
#   nvidia-smi.txt        current GPU snapshot + per-GPU CSV row
#   docker-info.txt       docker version/info + /etc/docker/daemon.json
#   docker-inspect-<id>   per-matched-container inspect JSON (env redacted)
#   docker-<id>.log       per-matched-container logs (since 24h, cap 5 MB)
#   vllm-proc.txt         if no docker: PID, cmdline, mounts, fd snapshot
#   journal-24h.log       systemd journal last 24h (cap 5 MB)
#   extra-<basename>      contents of --log files (cap 5 MB each)
#
# Sensitive env values (keys matching TOKEN/SECRET/KEY/PASSWORD/AUTH) are
# redacted before upload. We never collect /etc/shadow, /root/.ssh, or any
# file outside the paths listed above.

set -uo pipefail

# ---------------- parse args ----------------
TOKEN=""
IMAGE_FILTER=""
CONTAINER_FILTER=""
declare -a EXTRA_LOGS=()
ENDPOINT_BASE="${KAITAKU_BASE_URL:-https://kaitaku.ai}"

while [ $# -gt 0 ]; do
  case "$1" in
    --image)     IMAGE_FILTER="$2"; shift 2;;
    --container) CONTAINER_FILTER="$2"; shift 2;;
    --log)       EXTRA_LOGS+=("$2"); shift 2;;
    --base)      ENDPOINT_BASE="$2"; shift 2;;
    -h|--help)
      grep '^#' "$0" | sed 's/^# \{0,1\}//'
      exit 0;;
    -*)
      echo "unknown flag: $1" >&2; exit 2;;
    *)
      if [ -z "$TOKEN" ]; then TOKEN="$1"; else echo "extra positional: $1" >&2; exit 2; fi
      shift;;
  esac
done

if [ -z "$TOKEN" ]; then
  echo "usage: bash <(curl -sSL kaitaku.ai/report.sh) <TOKEN> [--image NAME:TAG] [--log /path]" >&2
  exit 2
fi

ENDPOINT_BASE="${ENDPOINT_BASE%/}"
UPLOAD_URL="${ENDPOINT_BASE}/api/feedback/${TOKEN}/bundle"

# ---------------- staging ----------------
WORK="$(mktemp -d -t kaitaku-report.XXXXXX)"
trap 'rm -rf "$WORK"' EXIT
cd "$WORK" || exit 1

MANIFEST="$WORK/manifest.txt"
collected=()
skipped=()

# ---------------- environment detection ----------------
ENV_KIND="bare-metal-no-docker"   # default fallback
if [ -f /.dockerenv ] || grep -qE '/(docker|containerd|kubepods|lxc)/' /proc/1/cgroup 2>/dev/null; then
  ENV_KIND="inside-container"
elif command -v docker >/dev/null 2>&1 && docker info >/dev/null 2>&1; then
  ENV_KIND="bare-metal-with-docker"
fi

HOSTNAME_REAL="$(hostname 2>/dev/null || echo unknown)"

# Our mlnode/vllm image streams stdout straight to Docker on the host — there
# is no log file inside the container. If Docker is unreachable AND the caller
# didn't hand us an explicit `--log /path/to/log`, this bundle would only
# carry metadata (sysinfo + nvidia-smi). That's misleading, so refuse early
# and tell the user where to actually get the logs.
if [ "$ENV_KIND" != "bare-metal-with-docker" ] && [ ${#EXTRA_LOGS[@]} -eq 0 ]; then
  cat >&2 <<EOF
[!] $ENV_KIND detected — vllm logs are not collectable from here.

In our image, vllm's stdout goes to Docker on the HOST, not to any file
inside the container. To get real logs, either:

  1. Run this script ON THE HOST that owns the container, e.g.:
     curl -sSL ${ENDPOINT_BASE}/report.sh | bash -s -- $TOKEN \\
       --image <image-name:tag>

  2. Grab logs from your provider's web console (Vast.ai / Runpod / Lambda)
     and upload them via the "Upload a file" tab in the feedback modal.

  3. If you already have a log file here, pass it explicitly:
     bash $0 $TOKEN --log /path/to/your/vllm.log

EOF
  exit 5
fi

{
  echo "=== kaitaku diagnostic bundle ==="
  echo "generated:        $(date -u +%Y-%m-%dT%H:%M:%SZ)"
  echo "host:             $HOSTNAME_REAL"
  echo "token:            ${TOKEN:0:8}…"
  echo "endpoint:         $UPLOAD_URL"
  echo "env kind:         $ENV_KIND"
  echo "filter --image:   ${IMAGE_FILTER:-(none)}"
  echo "filter --container: ${CONTAINER_FILTER:-(none)}"
  echo "user --log paths: ${EXTRA_LOGS[*]:-(none)}"
  echo
} > "$MANIFEST"

# ---------------- sysinfo (always) ----------------
{
  echo "## uname"
  uname -a 2>&1 || echo "uname failed"
  echo
  echo "## /etc/os-release"
  cat /etc/os-release 2>&1 || echo "no /etc/os-release"
  echo
  echo "## memory"
  (free -h 2>/dev/null || vm_stat 2>/dev/null) || echo "no free/vm_stat"
  echo
  echo "## disk (only the busiest mounts)"
  df -h 2>&1 | head -20 || echo "df failed"
  echo
  echo "## env (filtered, no secrets)"
  env | grep -Ei '^(CUDA|LD_LIBRARY|NVIDIA|VLLM|HF_|DOCKER|HOSTNAME|USER|PATH|TZ)=' \
      | sed -E 's/^([^=]*(TOKEN|KEY|SECRET|PASS|AUTH))=.*/\1=<redacted>/Ig' \
      || true
} > sysinfo.txt 2>&1
collected+=("sysinfo.txt")

# ---------------- nvidia-smi (works in all 3 envs) ----------------
if command -v nvidia-smi >/dev/null 2>&1; then
  nvidia-smi > nvidia-smi.txt 2>&1 || true
  {
    echo
    echo "## per-GPU CSV"
    nvidia-smi --query-gpu=index,name,driver_version,memory.total,memory.used,utilization.gpu,temperature.gpu \
               --format=csv 2>&1 || true
    echo
    echo "## per-GPU UUID list (matches container GPU device requests)"
    nvidia-smi -L 2>&1 || true
  } >> nvidia-smi.txt
  collected+=("nvidia-smi.txt")
else
  skipped+=("nvidia-smi (binary not found)")
fi

# ---------------- env-specific collectors ----------------
sanitise_inspect_to() {
  # $1: src JSON path, $2: dest path
  # Redact Config.Env values where key matches TOKEN/SECRET/KEY/PASSWORD/AUTH.
  # python3 is available in nearly every box we care about; if not, copy raw.
  if command -v python3 >/dev/null 2>&1; then
    python3 - "$1" "$2" <<'PY' 2>/dev/null || cp "$1" "$2"
import json, re, sys
src, dst = sys.argv[1], sys.argv[2]
pat = re.compile(r'(?i)(TOKEN|SECRET|KEY|PASSWORD|PASS|AUTH)')
try:
    data = json.load(open(src))
except Exception:
    open(dst, 'wb').write(open(src, 'rb').read()); raise SystemExit(0)
def redact(env):
    out = []
    for kv in env or []:
        if '=' not in kv: out.append(kv); continue
        k, v = kv.split('=', 1)
        out.append(f"{k}=<redacted>" if pat.search(k) else kv)
    return out
for c in data if isinstance(data, list) else [data]:
    try:
        c['Config']['Env'] = redact(c['Config'].get('Env'))
    except Exception:
        pass
json.dump(data, open(dst, 'w'), indent=2)
PY
  else
    cp "$1" "$2"
  fi
}

cap_to_5mb() {
  # In-place truncate $1 to its last 5 MB if larger.
  local f="$1"
  local size
  size=$(wc -c < "$f" 2>/dev/null || echo 0)
  if [ "$size" -gt $((5 * 1024 * 1024)) ]; then
    tail -c $((5 * 1024 * 1024)) "$f" > "$f.tail"
    mv "$f.tail" "$f"
  fi
}

if [ "$ENV_KIND" = "bare-metal-with-docker" ]; then
  {
    echo "## docker version"
    docker version 2>&1 || true
    echo
    echo "## docker info (head)"
    docker info 2>&1 | head -40 || true
    echo
    echo "## /etc/docker/daemon.json"
    cat /etc/docker/daemon.json 2>/dev/null || echo "no daemon.json"
    echo
    echo "## docker ps -a"
    docker ps -a --format 'table {{.ID}}\t{{.Image}}\t{{.Names}}\t{{.Status}}' 2>&1 || true
  } > docker-info.txt
  collected+=("docker-info.txt")

  # ---- pick target containers ----
  declare -a CIDS=()
  if [ -n "$CONTAINER_FILTER" ]; then
    cid="$(docker ps -aq --filter "name=$CONTAINER_FILTER" 2>/dev/null | head -1)"
    [ -z "$cid" ] && cid="$(docker inspect --format '{{.Id}}' "$CONTAINER_FILTER" 2>/dev/null || true)"
    [ -n "$cid" ] && CIDS+=("$cid") || skipped+=("--container $CONTAINER_FILTER did not match anything")
  elif [ -n "$IMAGE_FILTER" ]; then
    # Three-tier match:
    #   1. Exact `ancestor=<image:tag>`
    #   2. Substring of the full `<id> <image>` line (handles tag variants like
    #      `<base>` vs `<base>-kimi-1`)
    #   3. Same image NAME (before the colon), any tag — so a report about
    #      `mlnode-full:b300-k5` still picks up `mlnode-full:b300-k5-kimi-1`,
    #      `mlnode-full:b200-k5-kimi-1`, etc.
    while read -r cid; do
      [ -n "$cid" ] && CIDS+=("$cid")
    done < <(docker ps -aq --filter "ancestor=$IMAGE_FILTER" 2>/dev/null)
    if [ ${#CIDS[@]} -eq 0 ]; then
      while read -r cid; do
        [ -n "$cid" ] && CIDS+=("$cid")
      done < <(docker ps -a --no-trunc --format '{{.ID}} {{.Image}}' 2>/dev/null \
                 | grep -F "$IMAGE_FILTER" | awk '{print $1}' | head -8)
    fi
    if [ ${#CIDS[@]} -eq 0 ]; then
      IMAGE_NAME_ONLY="${IMAGE_FILTER%%:*}"
      if [ -n "$IMAGE_NAME_ONLY" ] && [ "$IMAGE_NAME_ONLY" != "$IMAGE_FILTER" ]; then
        while read -r cid; do
          [ -n "$cid" ] && CIDS+=("$cid")
        done < <(docker ps -a --no-trunc --format '{{.ID}} {{.Image}}' 2>/dev/null \
                   | awk -v name="$IMAGE_NAME_ONLY" '$2 ~ "^"name":" {print $1}' | head -8)
        [ ${#CIDS[@]} -gt 0 ] && skipped+=("--image $IMAGE_FILTER not found verbatim; matched by name $IMAGE_NAME_ONLY:*")
      fi
    fi
    [ ${#CIDS[@]} -eq 0 ] && skipped+=("--image $IMAGE_FILTER did not match any container")
  fi
  # No filter or no matches → fall back to all running containers (max 8)
  if [ ${#CIDS[@]} -eq 0 ]; then
    while read -r cid; do
      [ -n "$cid" ] && CIDS+=("$cid")
    done < <(docker ps -q 2>/dev/null | head -8)
  fi

  for cid in "${CIDS[@]}"; do
    short="${cid:0:12}"
    # Inspect (sanitise env)
    raw="$WORK/_inspect-${short}.raw.json"
    if docker inspect "$cid" > "$raw" 2>/dev/null; then
      sanitise_inspect_to "$raw" "docker-inspect-${short}.json"
      rm -f "$raw"
      collected+=("docker-inspect-${short}.json")
    fi
    # Logs (cap 5 MB)
    docker logs --since 24h --timestamps "$cid" > "docker-${short}.log" 2>&1 || true
    cap_to_5mb "docker-${short}.log"
    if [ -s "docker-${short}.log" ]; then
      collected+=("docker-${short}.log")
    else
      rm -f "docker-${short}.log"
    fi
  done

fi

# ---------------- journalctl (host fallback) ----------------
# Only meaningful on a bare-metal host. Inside a container we already exited
# above unless the caller provided --log files, in which case journalctl
# wouldn't help anyway (and usually isn't installed).
if [ "$ENV_KIND" = "bare-metal-with-docker" ] && command -v journalctl >/dev/null 2>&1; then
  journalctl --since "-24h" --no-pager > journal-24h.log 2>&1 || true
  cap_to_5mb journal-24h.log
  if [ -s journal-24h.log ]; then
    collected+=("journal-24h.log")
  else
    rm -f journal-24h.log
    skipped+=("journalctl (empty for last 24h)")
  fi
fi

# ---------------- explicit --log files ----------------
for src in "${EXTRA_LOGS[@]:-}"; do
  [ -z "$src" ] && continue
  if [ ! -r "$src" ]; then
    skipped+=("--log $src (not readable)")
    continue
  fi
  base="extra-$(basename "$src")"
  tail -c $((5 * 1024 * 1024)) "$src" > "$base" 2>/dev/null || true
  if [ -s "$base" ]; then
    collected+=("$base")
  else
    rm -f "$base"
  fi
done

# ---------------- manifest tail ----------------
{
  echo "## collected"
  for f in "${collected[@]}"; do echo "  - $f ($(wc -c < "$f") bytes)"; done
  echo
  echo "## skipped"
  for s in "${skipped[@]:-}"; do echo "  - $s"; done
} >> "$MANIFEST"

if [ ${#collected[@]} -eq 0 ]; then
  echo "[!] nothing useful to collect — refusing to upload an empty bundle." >&2
  cat "$MANIFEST" >&2
  exit 3
fi

# ---------------- bundle ----------------
tar czf bundle.tgz manifest.txt "${collected[@]}"
SIZE=$(wc -c < bundle.tgz | tr -d ' ')
echo "[+] bundle: ${SIZE} bytes, ${#collected[@]} files, env=$ENV_KIND"
cat "$MANIFEST"

# ---------------- upload ----------------
echo "[+] uploading to ${UPLOAD_URL} …"
HTTP_STATUS=$(curl -sS -o upload-response.json -w '%{http_code}' \
  -X POST \
  --data-binary @bundle.tgz \
  -H "Content-Type: application/octet-stream" \
  -H "X-Bundle-Sources: $(IFS=,; echo "${collected[*]}")" \
  -H "X-Bundle-Hostname: $HOSTNAME_REAL" \
  -H "X-Bundle-Env: $ENV_KIND" \
  -H "X-Bundle-Script: report.sh/2" \
  "${UPLOAD_URL}" || echo "000")

if [ "$HTTP_STATUS" = "200" ]; then
  echo "[✓] uploaded ok. response:"
  cat upload-response.json
  echo
  exit 0
fi
echo "[!] upload failed (HTTP $HTTP_STATUS). server said:" >&2
cat upload-response.json >&2 || true
echo >&2
exit 4