#!/usr/bin/env bash
# red-team-sweep.sh — Red-team every kindergarten seed sequentially.
# For each seed: inject into OptimizerV3.sol → run red-team.sh → restore → next.
# Usage: bash red-team-sweep.sh [timeout_per_candidate]
set -euo pipefail

REPO_ROOT="$(cd "$(dirname "$0")/../.." && pwd)"
SEEDS_DIR="$REPO_ROOT/tools/push3-evolution/seeds"
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
INJECT="$REPO_ROOT/tools/push3-transpiler/inject.sh"
ATTACKS_OUT="$REPO_ROOT/onchain/script/backtesting/attacks"
PROGRESS_FILE="/tmp/red-team-sweep-progress.json"
MEMORY_FILE="$REPO_ROOT/tmp/red-team-memory.jsonl"
CROSS_PATTERNS_FILE="$REPO_ROOT/tools/red-team/cross-patterns.jsonl"
SWEEP_TSV="/tmp/sweep-results.tsv"
OPT_SOL="$REPO_ROOT/onchain/src/OptimizerV3.sol"
TIMEOUT_PER="${1:-3600}"

log() { echo "[sweep $(date -u +%H:%M:%S)] $*"; }
die() { log "FATAL: $*" >&2; exit 1; }

[[ -f "$INJECT" ]] || die "inject.sh not found at $INJECT"

mkdir -p "$ATTACKS_OUT"
mkdir -p "$(dirname "$CROSS_PATTERNS_FILE")"

# Generate a unique sweep ID for traceability across runs
SWEEP_ID="sweep-$(date -u +%Y%m%d-%H%M%S)"
log "Sweep ID: $SWEEP_ID"

# Load progress
completed=()
if [[ -f "$PROGRESS_FILE" ]]; then
  while IFS= read -r line; do completed+=("$line"); done < <(jq -r '.completed[]' "$PROGRESS_FILE" 2>/dev/null || true)
fi
is_done() { for c in "${completed[@]+"${completed[@]}"}"; do [[ "$c" == "$1" ]] && return 0; done; return 1; }

# Collect named seeds only (skip run*_gen* pool entries)
seeds=()
for f in "$SEEDS_DIR"/*.push3; do
  [[ -f "$f" ]] || continue
  basename "$f" | grep -qE '^run[0-9]+_gen' && continue
  seeds+=("$f")
done
log "Found ${#seeds[@]} seeds. Timeout: ${TIMEOUT_PER}s each"
[[ ${#seeds[@]} -gt 0 ]] || die "No seeds found in $SEEDS_DIR"

# ── Smoke test: pick a random seed, inject + compile ──
SMOKE_IDX=$(( RANDOM % ${#seeds[@]} ))
SMOKE_SEED="${seeds[$SMOKE_IDX]}"
SMOKE_NAME=$(basename "$SMOKE_SEED" .push3)
log "Smoke test: $SMOKE_NAME"
cp "$OPT_SOL" "${OPT_SOL}.sweep-backup"
trap 'cp "${OPT_SOL}.sweep-backup" "$OPT_SOL" 2>/dev/null; rm -f "${OPT_SOL}.sweep-backup"' EXIT

bash "$INJECT" "$SMOKE_SEED" "$OPT_SOL" || die "Smoke test inject failed for $SMOKE_NAME"
cd "$REPO_ROOT/onchain" && forge build --silent 2>&1 || die "Smoke test compile failed for $SMOKE_NAME"
cp "${OPT_SOL}.sweep-backup" "$OPT_SOL"
log "Smoke test passed ✓"

# Write TSV header once (file persists across restarts; header only if new)
[[ ! -f "$SWEEP_TSV" ]] && \
  printf 'candidate\teth_before\teth_after\tpct_extracted\tstrategies_tried\tbest_attack\tstatus\n' > "$SWEEP_TSV"

# ── Main loop ──
for seed_file in "${seeds[@]}"; do
  seed_name=$(basename "$seed_file" .push3)
  is_done "$seed_name" && { log "SKIP $seed_name (done)"; continue; }

  log "=== RED-TEAM: $seed_name ==="

  # 1. Inject candidate into OptimizerV3.sol
  cp "${OPT_SOL}.sweep-backup" "$OPT_SOL"
  if ! bash "$INJECT" "$seed_file" "$OPT_SOL"; then
    log "SKIP $seed_name — inject failed"
    continue
  fi
  log "Injected into OptimizerV3.sol"

  # 1b. Extract optimizer profile from transpiler output (CI/AW/AS/DD constants)
  TRANSPILER_OUT="$REPO_ROOT/onchain/src/OptimizerV3Push3.sol"
  OPTIMIZER_PROFILE=$(python3 - "$TRANSPILER_OUT" <<'PYEOF'
import re, sys
try:
    with open(sys.argv[1]) as f:
        sol = f.read()
    ci_vals = set(re.findall(r'\br40\s*=\s*uint256\((\d+)\)', sol))
    aw_vals = set(re.findall(r'\br38\s*=\s*uint256\((\d+)\)', sol))
    as_vals = set(re.findall(r'\br39\s*=\s*uint256\((\d+)\)', sol))
    dd_vals = set(re.findall(r'\br37\s*=\s*uint256\((\d+)\)', sol))

    def fmt_pct(vals):
        pcts = sorted(set(round(int(v) * 100 / 1e18) for v in vals))
        return '/'.join(str(p) + '%' for p in pcts) if pcts else '?'

    def fmt_int(vals):
        ints = sorted(set(int(v) for v in vals))
        return '/'.join(str(v) for v in ints) if ints else '?'

    profile = f"CI={fmt_pct(ci_vals)}, AW={fmt_int(aw_vals)}, AS={fmt_pct(as_vals)}, DD={fmt_pct(dd_vals)}"
    # Adaptive: multiple constant branches, OR any register assigned from a variable
    has_var_assign = bool(re.search(r'\br(?:37|38|39|40)\s*=\s*uint256\s*\(\s*[a-zA-Z_]\w*\s*\)', sol))
    if len(ci_vals) > 1 or len(aw_vals) > 1 or len(as_vals) > 1 or len(dd_vals) > 1 or has_var_assign:
        profile += ", adaptive"
    print(profile)
except Exception as e:
    print(f"unknown (parse error: {e})", file=sys.stderr)
    print("unknown")
PYEOF
  )
  log "Optimizer profile: $OPTIMIZER_PROFILE"

  # 2. Clear stale attack file from previous candidate
  rm -f "$REPO_ROOT/tmp/red-team-attacks.jsonl"

  # 3. Run red-team.sh (handles bootstrap + compile + deploy + attack)
  log "Running red-team.sh (timeout: ${TIMEOUT_PER}s)..."
  set +e
  CLAUDE_TIMEOUT="$TIMEOUT_PER" CANDIDATE_NAME="$seed_name" OPTIMIZER_PROFILE="$OPTIMIZER_PROFILE" \
    timeout "$((TIMEOUT_PER + 120))" \
    bash "$SCRIPT_DIR/red-team.sh" 2>&1 | tee "/tmp/red-team-${seed_name}.log"
  RED_TEAM_EXIT="${PIPESTATUS[0]}"
  set -e

  # 4. Collect attacks
  if [[ -f "$REPO_ROOT/tmp/red-team-attacks.jsonl" ]]; then
    ATTACK_COUNT=$(wc -l < "$REPO_ROOT/tmp/red-team-attacks.jsonl")
    if [[ "$ATTACK_COUNT" -gt 0 ]]; then
      cp "$REPO_ROOT/tmp/red-team-attacks.jsonl" "$ATTACKS_OUT/sweep-${seed_name}.jsonl"
      log "Saved $ATTACK_COUNT attack(s)"
    fi
  fi

  # 4b. Write one TSV row to sweep-results.tsv
  # NOTE: intentionally runs before 4c (memory clear) so strategy data is still available.
  if [[ "$RED_TEAM_EXIT" -eq 0 ]]; then
    _sweep_status="safe"
  elif [[ "$RED_TEAM_EXIT" -eq 1 ]]; then
    _sweep_status="broken"
  elif [[ "$RED_TEAM_EXIT" -eq 124 ]]; then
    _sweep_status="timeout"
  else
    _sweep_status="crashed"
  fi
  set +e
  python3 - "/tmp/red-team-${seed_name}.log" "$MEMORY_FILE" "$seed_name" "$_sweep_status" "$SWEEP_TSV" <<'PYEOF'
import re, sys, json, os

log_file  = sys.argv[1]
mem_file  = sys.argv[2]
candidate = sys.argv[3]
status    = sys.argv[4]
tsv_file  = sys.argv[5]

# Parse eth_before (first occurrence = baseline) and eth_after (last occurrence = final state)
eth_before = ""
eth_after  = ""
try:
    with open(log_file) as f:
        for line in f:
            m = re.search(r'lm_eth_before\s*[=:]\s*(\d+)', line)
            if m and not eth_before:          # first occurrence wins
                eth_before = m.group(1)
            m = re.search(r'lm_eth_after\s*[=:]\s*(\d+)', line)
            if m:                             # last occurrence wins
                eth_after = m.group(1)
except Exception as e:
    print(f"  tsv: could not read log: {e}", file=sys.stderr)

# Parse strategies from the memory file (populated by extract_memory inside red-team.sh)
strategies_tried = 0
best_attack = "none"
try:
    if os.path.isfile(mem_file) and os.path.getsize(mem_file) > 0:
        with open(mem_file) as f:
            entries = [json.loads(l) for l in f if l.strip()]
        cand_entries = [e for e in entries if e.get("candidate") == candidate]
        strategies_tried = len(set(e["strategy"] for e in cand_entries if e.get("strategy")))
        best_delta = 0
        for e in cand_entries:
            if e.get("result") == "DECREASED" and e.get("delta_bps", 0) < best_delta:
                best_delta = e["delta_bps"]
                raw = e.get("strategy", "unknown")
                best_attack = re.sub(r"\s+", "_", raw.strip()).lower()[:50]
except Exception as e:
    print(f"  tsv: could not read memory: {e}", file=sys.stderr)

# Compute pct_extracted; use sentinel when ETH values are absent (crash/early-timeout)
if not eth_before and not eth_after:
    pct_extracted = ""
else:
    pct_extracted = "0.00"
    try:
        before = int(eth_before)
        after  = int(eth_after)
        if before > 0:
            extracted = max(0, before - after)
            pct_extracted = f"{extracted * 100 / before:.2f}"
    except Exception:
        pass

# Sanitise fields: strip tabs so the row is always valid TSV
def clean(s):
    return str(s).replace("\t", " ")

row = "\t".join([
    clean(candidate), clean(eth_before), clean(eth_after),
    clean(pct_extracted), clean(strategies_tried),
    clean(best_attack), clean(status),
])
with open(tsv_file, "a") as f:
    f.write(row + "\n")
print(f"  tsv: {status} | {pct_extracted}% extracted | {strategies_tried} strategies | best={best_attack}")
PYEOF
  _py_exit=$?
  set -e
  [[ $_py_exit -ne 0 ]] && log "WARNING: TSV row write failed (exit $_py_exit) — continuing"

  # 4c. Extract abstract patterns into cross-candidate file, then clear raw memory
  if [[ -f "$MEMORY_FILE" && -s "$MEMORY_FILE" ]]; then
    set +e
    _extract_out=$(python3 - "$MEMORY_FILE" "$CROSS_PATTERNS_FILE" "$SWEEP_ID" <<'PYEOF'
import json, sys

mem_file   = sys.argv[1]
cross_file = sys.argv[2]
sweep_id   = sys.argv[3] if len(sys.argv) > 3 else "unknown"

new_entries = []
with open(mem_file) as f:
    for line in f:
        line = line.strip()
        if line:
            try:
                new_entries.append(json.loads(line))
            except Exception:
                pass

if not new_entries:
    print("No memory entries to extract")
    sys.exit(0)

# Load existing (pattern, candidate, result) keys for deduplication
existing_keys = set()
try:
    with open(cross_file) as f:
        for line in f:
            line = line.strip()
            if line:
                try:
                    e = json.loads(line)
                    existing_keys.add((e.get("pattern", ""), e.get("candidate", ""), e.get("result", "")))
                except Exception:
                    pass
except FileNotFoundError:
    pass

appended = 0
skipped = 0
with open(cross_file, 'a') as f:
    for e in new_entries:
        key = (e.get("pattern", ""), e.get("candidate", ""), e.get("result", ""))
        if key in existing_keys:
            skipped += 1
            continue
        e["sweep_id"] = sweep_id
        existing_keys.add(key)  # prevent intra-batch duplicates
        f.write(json.dumps(e) + '\n')
        appended += 1

total = appended + skipped
print(f"Extracted {appended} new entr{'y' if appended==1 else 'ies'} ({skipped} duplicate{'s' if skipped!=1 else ''} skipped) to cross-patterns file")
PYEOF
    )
    _py_exit=$?
    set -e
    [[ -n "$_extract_out" ]] && log "$_extract_out"
    [[ $_py_exit -ne 0 ]] && log "WARNING: cross-pattern extraction failed (exit $_py_exit) — continuing"
  fi
  # Always clear raw memory so the next candidate starts with a fresh tactical state
  if [[ -f "$MEMORY_FILE" ]]; then
    > "$MEMORY_FILE"
    log "Cleared raw memory for next candidate"
  fi

  # 5. Save progress
  completed+=("$seed_name")
  jq -n --argjson arr "$(printf '%s\n' "${completed[@]}" | jq -R . | jq -s .)" \
    '{completed: $arr, last_updated: now | todate}' > "$PROGRESS_FILE"
  log "DONE $seed_name"

  # 6. Teardown — poll until all containers have exited (no fixed sleep)
  cd "$REPO_ROOT" && docker compose down -v 2>/dev/null || true
  _deadline=$(( $(date +%s) + 30 ))
  while [[ -n "$(docker compose ps --quiet 2>/dev/null)" ]]; do
    if [[ $(date +%s) -ge $_deadline ]]; then
      log "WARNING: containers still present after 30s — proceeding anyway"
      break
    fi
    sleep 1
  done
done

# Restore original
cp "${OPT_SOL}.sweep-backup" "$OPT_SOL"
log "=== SWEEP COMPLETE: ${#completed[@]} / ${#seeds[@]} ==="