#!/usr/bin/env bash
# audit-wiring.sh — enforce "every kernel/slot/head/action fully wired"
# per SP20 §0.2.
#
# Walks all four manifests:
#   * kernels.txt → every entry has a registration in build.rs KERNELS
#                   AND a `.launch_builder(&self.<name>_fn)` site
#                   in src/trainer/integrated.rs
#   * slots.txt   → every entry has a write site (seed in rl_isv_write
#                   OR write in a .cu kernel) AND a read site
#                   (read in a .cu kernel)
#   * heads.txt   → every entry has src/heads/<name>.rs AND forward
#                   call AND backward call in the trainer
#   * actions.txt → every entry has a branch in
#                   cuda/actions_to_market_targets.cu OR a handler
#                   in an override kernel
#
# Exit codes:
#   0  — zero violations
#   1  — one or more violations
#   2  — usage error
# Note: NOT using `set -e` — many grep|head|wc pipelines naturally
# return non-zero (SIGPIPE on head consuming early, grep no-match).
# We handle exit codes explicitly via `|| true` and `|| echo 0` on
# the suspect pipelines, and bookkeep `violations` for the final
# exit code.
set -uo pipefail

ROOT=$(git rev-parse --show-toplevel)
M="$ROOT/scripts/audit-manifest"
BUILD_RS="$ROOT/crates/ml-alpha/build.rs"
TRAINER="$ROOT/crates/ml-alpha/src/trainer/integrated.rs"
ISV_SLOTS="$ROOT/crates/ml-alpha/src/rl/isv_slots.rs"
COMMON_RS="$ROOT/crates/ml-alpha/src/rl/common.rs"
ACTIONS_KERNEL="$ROOT/crates/ml-alpha/cuda/actions_to_market_targets.cu"
HEADS_DIR="$ROOT/crates/ml-alpha/src/heads"
CUDA_DIR="$ROOT/crates/ml-alpha/cuda"

violations=0

read_manifest() {
  local file="$1"
  if [[ -f "$file" ]]; then
    grep -vE '^\s*(#|$)' "$file" | awk '{print $1}'
  fi
}

# ─── 1. KERNELS ─────────────────────────────────────────────
echo "audit-wiring: KERNELS"
while IFS= read -r kernel; do
  [[ -z "$kernel" ]] && continue
  cu_file="$CUDA_DIR/${kernel}.cu"
  if [[ ! -f "$cu_file" ]]; then
    echo "  VIOLATION kernel $kernel: source file missing ($cu_file)"
    violations=$((violations + 1))
    continue
  fi
  # build.rs registration
  if ! grep -qE "\"${kernel}\"" "$BUILD_RS"; then
    echo "  VIOLATION kernel $kernel: not registered in build.rs KERNELS"
    violations=$((violations + 1))
  fi
  # trainer field
  if ! grep -qE "${kernel}_fn[: ,]" "$TRAINER"; then
    echo "  VIOLATION kernel $kernel: no <name>_fn field in trainer"
    violations=$((violations + 1))
  fi
  # trainer launch site
  if ! grep -qE "launch_builder\(&self\.${kernel}_fn\)" "$TRAINER"; then
    echo "  VIOLATION kernel $kernel: no launch site in trainer"
    violations=$((violations + 1))
  fi
done < <(read_manifest "$M/kernels.txt")

# ─── 2. SLOTS ───────────────────────────────────────────────
echo "audit-wiring: SLOTS"
while IFS= read -r slot; do
  [[ -z "$slot" ]] && continue
  # slot defined in isv_slots.rs?
  if ! grep -qE "pub const ${slot}\s*:\s*usize" "$ISV_SLOTS"; then
    echo "  VIOLATION slot $slot: no `pub const` definition in isv_slots.rs"
    violations=$((violations + 1))
    continue
  fi
  # Extract numeric slot index for cross-kernel grep
  slot_idx=$(awk -v name="$slot" '
    $0 ~ "pub const " name "[[:space:]]*:[[:space:]]*usize" {
      if (match($0, /=[[:space:]]*[0-9]+/)) {
        v = substr($0, RSTART, RLENGTH)
        gsub(/[^0-9]/, "", v)
        print v; exit
      }
    }' "$ISV_SLOTS")
  [[ -z "$slot_idx" ]] && slot_idx="-1"
  # producer: seeded in rl_isv_write list OR written by a .cu kernel
  producer_seed=0
  producer_kernel=0
  if grep -qE "isv_slots::${slot}" "$TRAINER"; then
    producer_seed=1
  fi
  # kernel writes via isv[INDEX] = ...  or isv[<num>] = ...
  if grep -qrE "isv\[(${slot}|${slot_idx})\][[:space:]]*=" "$CUDA_DIR" 2>/dev/null; then
    producer_kernel=1
  fi
  if [[ $producer_seed -eq 0 && $producer_kernel -eq 0 ]]; then
    echo "  VIOLATION slot $slot (idx $slot_idx): no producer (not seeded in rl_isv_write, not written by any kernel)"
    violations=$((violations + 1))
  fi
  # consumer: read by some .cu (isv[INDEX] in non-assignment context)
  # Heuristic: count distinct .cu files containing the slot reference
  # in a context that's NOT immediately followed by `=` (assignment).
  # Pattern `isv[X][^=]` matches isv[X] where next char isn't `=`.
  read_only_files=$(grep -lrE "isv\[(${slot}|${slot_idx})\][^=]" "$CUDA_DIR" 2>/dev/null | wc -l)
  if [[ $read_only_files -eq 0 ]]; then
    echo "  VIOLATION slot $slot (idx $slot_idx): no consumer (no kernel reads it in non-assignment context)"
    violations=$((violations + 1))
  fi
done < <(read_manifest "$M/slots.txt")

# ─── 3. HEADS ───────────────────────────────────────────────
echo "audit-wiring: HEADS"
while IFS= read -r head; do
  [[ -z "$head" ]] && continue
  head_rs="$HEADS_DIR/${head}.rs"
  if [[ ! -f "$head_rs" ]]; then
    echo "  VIOLATION head $head: module file missing ($head_rs)"
    violations=$((violations + 1))
    continue
  fi
  if ! grep -qE "fn forward" "$head_rs"; then
    echo "  VIOLATION head $head: no fn forward in module"
    violations=$((violations + 1))
  fi
  if ! grep -qE "fn backward" "$head_rs"; then
    echo "  VIOLATION head $head: no fn backward in module"
    violations=$((violations + 1))
  fi
  # Adam step + LR controller — grep for the head's forward call site
  # in the trainer (loose heuristic).
  if ! grep -qE "${head}_head" "$TRAINER"; then
    echo "  VIOLATION head $head: no ${head}_head field or call in trainer"
    violations=$((violations + 1))
  fi
done < <(read_manifest "$M/heads.txt")

# ─── 4. ACTIONS ─────────────────────────────────────────────
echo "audit-wiring: ACTIONS"
while IFS= read -r action; do
  [[ -z "$action" ]] && continue
  # Action enum entry
  if ! grep -qE "${action}\s*=\s*[0-9]+" "$COMMON_RS"; then
    echo "  VIOLATION action $action: not in Action enum in common.rs"
    violations=$((violations + 1))
    continue
  fi
  # Get action index
  action_idx=$(awk -v name="$action" '
    $0 ~ name "[[:space:]]*=[[:space:]]*[0-9]+" {
      if (match($0, /=[[:space:]]*[0-9]+/)) {
        v = substr($0, RSTART, RLENGTH)
        gsub(/[^0-9]/, "", v)
        print v; exit
      }
    }' "$COMMON_RS")
  [[ -z "$action_idx" ]] && action_idx="-1"
  # Handler patterns (any of):
  #   * Literal index branch in any .cu: `action == <idx>` or `actions[..] == <idx>`
  #   * Named-constant compare: `action == ACTION_<UPPER_SNAKE>` where
  #     UPPER_SNAKE is the CamelCase enum name → UPPER_SNAKE_CASE
  #   * Disequality (`!=`) is also a handler (a guard checking for the action)
  upper_snake=$(echo "$action" | sed -E 's/([A-Z])/_\1/g; s/^_//' | tr '[:lower:]' '[:upper:]')
  named_const="ACTION_${upper_snake}"
  pat_literal="(action|actions\\[[^]]*\\])[[:space:]]*[!=]=[[:space:]]*${action_idx}\\b"
  pat_named="(action|actions\\[[^]]*\\])[[:space:]]*[!=]=[[:space:]]*${named_const}\\b"

  handler_count=$(grep -lrE "(${pat_literal})|(${pat_named})" "$CUDA_DIR" 2>/dev/null | wc -l)
  if [[ "$handler_count" -eq 0 ]]; then
    echo "  VIOLATION action $action (idx $action_idx): no handler in any kernel"
    echo "    (searched for: \`action == $action_idx\` or \`action == $named_const\`)"
    violations=$((violations + 1))
  fi
done < <(read_manifest "$M/actions.txt")

echo
echo "audit-wiring: $violations violation(s)"

if [[ $violations -gt 0 ]]; then
  echo
  echo "Per SP20 §0.2: every kernel/slot/head/action MUST have producer + consumer in same commit."
  exit 1
fi

echo "audit-wiring: PASS"
exit 0