fizz

#!/bin/bash

# fail on error
set -e

SCRIPT_DIR="$(dirname "$(readlink -f "${BASH_SOURCE[0]:-$0}")")"
WORKING_DIR="$(pwd)"
SCRIPT_ENV_FILE="$SCRIPT_DIR/fizz.env"

# load envs
if [[ -f $SCRIPT_ENV_FILE ]]; then
  source $SCRIPT_ENV_FILE
fi

temp_output=$(mktemp)

MBT_GEN_ZIP="${MBT_GEN_ZIP:-$SCRIPT_DIR/bazel-bin/mbt/generator/generator_bin.zip}"

# --- SUBCOMMAND HANDLING ---
if [[ $# -gt 0 ]]; then
  case "$1" in
    mbt-scaffold)
      shift  # remove subcommand from $@
      "python" "$MBT_GEN_ZIP" "$@" > "$temp_output" || {
        echo "Error: mbt-scaffold command failed"
        echo "Logs at $temp_output"
        exit 1
      }
      exit 0
      ;;

    install-skills)
      shift  # remove subcommand from $@

      _FIZZ_REPO="fizzbee-io/fizzbee"
      _FIZZ_BRANCH="main"
      _FIZZ_RAW="https://raw.githubusercontent.com/${_FIZZ_REPO}/${_FIZZ_BRANCH}"
      _FIZZ_SKILLS_DIR="${HOME}/.claude/skills"
      _FIZZ_DOCS_DIR="${_FIZZ_SKILLS_DIR}/fizzbee-docs"
      _FIZZ_SKILLS=(fizz-spec fizz-check fizz-debug fizz-mbt)
      _FIZZ_DOCS=(
        "examples/references/LANGUAGE_REFERENCE.md"
        "examples/references/GOTCHAS.md"
        "examples/references/PERFORMANCE_GUIDE.md"
        "examples/references/VERIFICATION_GUIDE.md"
        "examples/references/README.md"
      )
      _FIZZ_EXAMPLE_NAMES=(
        "01-counter.fizz"
        "09-assertions.fizz"
        "11-roles.fizz"
        "13-two-phase-commit.fizz"
        "14-fault-injection.fizz"
        "16-symmetry.fizz"
      )
      _FIZZ_EXAMPLE_PATHS=(
        "examples/references/01-02-atomic-action/Counter.fizz"
        "examples/references/09-01-always-assertion/Counter.fizz"
        "examples/references/11-01-simple-role/Account.fizz"
        "examples/references/13-02-01-two-phase-commit/TwoPhaseCommit.fizz"
        "examples/references/14-01-crash-on-yield/CrashOnYield.fizz"
        "examples/references/16-05-nominal-symmetry/NominalSymmetry.fizz"
      )

      if command -v curl &>/dev/null; then
        _DL="curl -fsSL"
      elif command -v wget &>/dev/null; then
        _DL="wget -qO-"
      else
        echo "Error: curl or wget is required to install skills" >&2
        exit 1
      fi

      _do_skills_check() {
        echo "Would install to: ${_FIZZ_SKILLS_DIR}/"
        echo ""
        echo "Skills:"
        for skill in "${_FIZZ_SKILLS[@]}"; do
          if [[ -f "${_FIZZ_SKILLS_DIR}/${skill}/SKILL.md" ]]; then
            echo "  ↻ ${skill}  (already installed, would update)"
          else
            echo "  + ${skill}  (new)"
          fi
        done
        echo ""
        echo "Reference docs (${_FIZZ_DOCS_DIR}/):"
        for doc in "${_FIZZ_DOCS[@]}"; do
          name=$(basename "$doc")
          if [[ -f "${_FIZZ_DOCS_DIR}/${name}" ]]; then
            echo "  ↻ ${name}  (already installed, would update)"
          else
            echo "  + ${name}  (new)"
          fi
        done
      }

      _do_skills_remove() {
        echo "Removing FizzBee Claude skills and docs..."
        for skill in "${_FIZZ_SKILLS[@]}"; do
          if [[ -d "${_FIZZ_SKILLS_DIR}/${skill}" ]]; then
            rm -rf "${_FIZZ_SKILLS_DIR:?}/${skill}"
            echo "  ✗ ${skill}"
          else
            echo "  - ${skill}  (not installed)"
          fi
        done
        if [[ -d "${_FIZZ_DOCS_DIR}" ]]; then
          rm -rf "${_FIZZ_DOCS_DIR:?}"
          echo "  ✗ fizzbee-docs/"
        else
          echo "  - fizzbee-docs/  (not installed)"
        fi
        echo ""
        echo "Done."
      }

      _do_skills_install() {
        mkdir -p "${_FIZZ_SKILLS_DIR}" "${_FIZZ_DOCS_DIR}" "${_FIZZ_DOCS_DIR}/examples"

        echo "Installing FizzBee Claude skills to ${_FIZZ_SKILLS_DIR}/"
        echo ""
        for skill in "${_FIZZ_SKILLS[@]}"; do
          mkdir -p "${_FIZZ_SKILLS_DIR}/${skill}"
          $_DL "${_FIZZ_RAW}/.claude/skills/${skill}/SKILL.md" \
            > "${_FIZZ_SKILLS_DIR}/${skill}/SKILL.md"
          echo "  ✓ ${skill}"
        done

        echo ""
        echo "Installing reference docs to ${_FIZZ_DOCS_DIR}/"
        echo ""
        for doc in "${_FIZZ_DOCS[@]}"; do
          name=$(basename "$doc")
          $_DL "${_FIZZ_RAW}/${doc}" > "${_FIZZ_DOCS_DIR}/${name}"
          echo "  ✓ ${name}"
        done

        echo ""
        echo "Installing example specs to ${_FIZZ_DOCS_DIR}/examples/"
        echo ""
        for i in "${!_FIZZ_EXAMPLE_NAMES[@]}"; do
          $_DL "${_FIZZ_RAW}/${_FIZZ_EXAMPLE_PATHS[$i]}" \
            > "${_FIZZ_DOCS_DIR}/examples/${_FIZZ_EXAMPLE_NAMES[$i]}"
          echo "  ✓ ${_FIZZ_EXAMPLE_NAMES[$i]}"
        done

        echo ""
        echo "Done. Skills are now available in all Claude Code sessions."
        echo ""
        echo "To use:"
        echo "  - Claude Code will auto-invoke the right skill based on context"
        echo "  - Or invoke manually: /fizz-spec  /fizz-check  /fizz-debug  /fizz-mbt"
        echo ""
        echo "To update after upgrading fizzbee, run: fizz install-skills"
      }

      case "${1:-}" in
        --check)   _do_skills_check ;;
        --remove)  _do_skills_remove ;;
        *)         _do_skills_install ;;
      esac
      exit 0
      ;;
  esac
fi

usage() {
  echo "Usage:"
  echo "  $0 install-skills [--check | --remove]"
  echo "  $0 mbt-scaffold [options] filename"
  echo "  $0 [-x|--simulation] [--test] [--seed int64Number] [--max_runs intNumber] [--simulation_first_traces intNumber] [--exploration_strategy dfs] [--trace-file tracefile | --trace tracestring] [--preinit-hook-file hookfile | --preinit-hook hookstring] [--copy-ast] [--no-copy-ast] [--output-dir directory] [--parallel intNumber] [experimental] filename"
  echo ""
  echo "  --parallel N"
  echo "    Run N simulation processes in parallel (only with -x and no --seed)."
  echo "    Workers split --max_runs evenly. Each worker writes to its own"
  echo "    sub-directory; the first failing worker's output is shown."
  echo "    Default=1 (sequential)."
  echo ""
  echo "  Experimental flags (advanced; off by default):"
  echo "    --experimental_processed_queue"
  echo "        Queue holds processed (yield-point) nodes instead of unprocessed"
  echo "        action-starts. Dedupes successors before they enter the queue,"
  echo "        reducing peak queue memory. Works with BFS, DFS, or random."
  echo "    --experimental_no_graph"
  echo "        Drops the in-memory state graph after each yield-point is"
  echo "        expanded. Major RSS reduction; refuses specs with liveness"
  echo "        assertions; disables crash-on-yield simulation. Auto-enables"
  echo "        --experimental_processed_queue."
  echo "    --experimental_no_state_returns"
  echo "        Omits Process.Returns from HashCode, JSON state, and dot-file"
  echo "        state labels. Returns remain available on the per-transition"
  echo "        Link.Returns field. Use to verify downstream consumers (graph,"
  echo "        MBT, explorer) have migrated to reading from links."
  exit 1
}

# Initialize variables
simulation=false
seed=0
max_runs=0
simulation_first_traces=0
exploration_strategy=bfs
test=false
copy_ast=true
output_dir=""
trace_file=""
trace_string=""
trace_extend=0
preinit_hook_file=""
preinit_hook_string=""
experimental_processed_queue=false
experimental_no_graph=false
experimental_no_state_returns=false
parallel=1

# Parse options
while [[ "$1" =~ ^- ]]; do
  case $1 in
    -x | --simulation )
      simulation=true
      shift
      ;;
    --test )
      test=true
      shift
      ;;
    --seed )
      if [[ -n "$2" ]] && [[ "$2" =~ ^[0-9]+$ ]]; then
        seed="$2"
        shift 2
      else
        echo "Error: --seed requires a numeric value." 1>&2
        usage
      fi
      ;;
    --max_runs )
      if [[ -n "$2" ]] && [[ "$2" =~ ^[0-9]+$ ]]; then
        max_runs="$2"
        shift 2
      else
        echo "Error: --max_runs requires a numeric value." 1>&2
        usage
      fi
      ;;
    --simulation_first_traces )
      if [[ -n "$2" ]] && [[ "$2" =~ ^[0-9]+$ ]]; then
        simulation_first_traces="$2"
        shift 2
      else
        echo "Error: --simulation_first_traces requires a numeric value." 1>&2
        usage
      fi
      ;;
    --exploration_strategy )
      if [[ -n "$2" ]]; then
        exploration_strategy="$2"
        shift 2
      else
        echo "Error: --exploration_strategy requires a string value one of [bfs, dfs, random]." 1>&2
        usage
      fi
      ;;
    --trace-file )
      if [[ -n "$2" ]]; then
        trace_file="$2"
        shift 2
      else
        echo "Error: --trace-file requires a file path." 1>&2
        usage
      fi
      ;;
    --trace )
      if [[ -n "$2" ]]; then
        trace_string="$2"
        shift 2
      else
        echo "Error: --trace requires a string." 1>&2
        usage
      fi
      ;;
    --trace-extend )
      if [[ -n "$2" ]] && [[ "$2" =~ ^[0-9]+$ ]]; then
        trace_extend="$2"
        shift 2
      else
        echo "Error: --trace-extend requires a numeric value." 1>&2
        usage
      fi
      ;;
    --preinit-hook-file )
      if [[ -n "$2" ]]; then
        preinit_hook_file="$2"
        shift 2
      else
        echo "Error: --preinit-hook-file requires a file path." 1>&2
        usage
      fi
      ;;
    --preinit-hook )
      if [[ -n "$2" ]]; then
        preinit_hook_string="$2"
        shift 2
      else
        echo "Error: --preinit-hook requires a string." 1>&2
        usage
      fi
      ;;
    --copy-ast )
      copy_ast=true
      shift
      ;;
    --no-copy-ast )
      copy_ast=false
      shift
      ;;
    --output-dir )
      if [[ -n "$2" ]]; then
        output_dir="$2"
        shift 2
      else
        echo "Error: --output-dir requires a directory path." 1>&2
        usage
      fi
      ;;
    --internal_profile )
      internal_profile=true
      shift
      ;;
    --experimental_processed_queue )
      experimental_processed_queue=true
      shift
      ;;
    --experimental_no_graph )
      experimental_no_graph=true
      shift
      ;;
    --experimental_no_state_returns )
      experimental_no_state_returns=true
      shift
      ;;
    --parallel )
      if [[ -n "$2" ]] && [[ "$2" =~ ^[0-9]+$ ]]; then
        parallel="$2"
        shift 2
      else
        echo "Error: --parallel requires a positive integer value." 1>&2
        usage
      fi
      ;;
    -h | --help )
      usage
      ;;
    * )
      echo "Invalid option: $1" 1>&2
      usage
      ;;
  esac
done

# Check for the required positional argument
if [ -z "$1" ]; then
  echo "Error: filename is required" 1>&2
  usage
fi

input_filename=$1

# Example usage of the parsed options and arguments
if [ "$simulation" = true ]; then
  echo "Simulation mode is enabled"
fi

input_filename=$1

# Set the paths for the binaries relative to the script directory
PARSER_BIN="${PARSER_BIN:-$SCRIPT_DIR/bazel-bin/parser/parser_bin}"
FIZZBEE_BIN="${FIZZBEE_BIN:-$SCRIPT_DIR/bazel-bin/fizzbee_/fizzbee}"

# Call the first binary and redirect output to a temporary file
if ! "$PARSER_BIN" "$input_filename" > "$temp_output"; then
    echo "Error: Compilation failed"
    echo "Logs at $temp_output"
    exit 1
fi

# Create the JSON filename by replacing the extension
json_filename="${input_filename%.*}.json"

echo "Model checking" $json_filename

# Prepare arguments for the Go binary
args=()
if [ "$simulation" = true ]; then
  args+=("--simulation")
fi
if [ "$test" = true ]; then
  args+=("--test")
fi
if [ "$internal_profile" = true ]; then
  args+=("--internal_profile")
fi
if [ "$experimental_processed_queue" = true ]; then
  args+=("--experimental_processed_queue")
fi
if [ "$experimental_no_graph" = true ]; then
  args+=("--experimental_no_graph")
fi
if [ "$experimental_no_state_returns" = true ]; then
  args+=("--experimental_no_state_returns")
fi
if [ "$seed" -ne 0 ]; then
  args+=("--seed" "$seed")
fi
if [ "$max_runs" -ne 0 ]; then
  args+=("--max_runs" "$max_runs")
fi
if [ "$simulation_first_traces" -ne 0 ]; then
  args+=("--simulation_first_traces" "$simulation_first_traces")
fi
if [ "$exploration_strategy" != "bfs" ]; then
  args+=("--exploration_strategy" "$exploration_strategy")
fi
if [ -n "$trace_file" ]; then
  args+=("--trace-file" "$trace_file")
fi
if [ -n "$trace_string" ]; then
  args+=("--trace" "$trace_string")
fi
if [ "$trace_extend" -ne 0 ]; then
  args+=("--trace-extend" "$trace_extend")
fi
if [ -n "$preinit_hook_file" ]; then
  args+=("--preinit-hook-file" "$preinit_hook_file")
fi
if [ -n "$preinit_hook_string" ]; then
  args+=("--preinit-hook" "$preinit_hook_string")
fi
if [ "$copy_ast" = true ]; then
  args+=("--copy-ast")
fi
if [ -n "$output_dir" ]; then
  args+=("--output-dir" "$output_dir")
fi

args+=("$json_filename")

# --- PARALLEL SIMULATION HANDLER -------------------------------------------
# Only fires when --simulation, no fixed seed, and --parallel > 1. Each
# worker is its own fizzbee subprocess — sidesteps all the in-process
# shared-state hazards (roleRefs, nextChannelId, etc.). Compatible with
# macOS bash 3.2 (uses only basic array and arithmetic features).
#
# Behavior:
# - Splits --max_runs evenly across workers (0 => unlimited per worker).
# - Each worker writes to its own out dir + log file.
# - Polling loop: on first worker that records a failure sentinel, kill
#   surviving workers. Up to ~500ms latency before survivors are killed.
# - Ctrl-C kills all live workers and exits with 130 (SIGINT convention).
# - On success: aggregate "PASSED: N total runs across W workers".
# - On failure: print the first failing worker's full log + its out dir.
if [ "$simulation" = true ] && [ "$parallel" -gt 1 ] && [ "$seed" -eq 0 ]; then
  # set -e is on at the top of the script; turn off inside this block so
  # expected non-zero exits (kill of an already-dead worker, grep with no
  # match, etc.) don't abort the script. We exit explicitly at the end.
  set +e

  # Pick a base output dir for this parallel run.
  if [ -n "$output_dir" ]; then
    parallel_base_dir="$output_dir/parallel_$(date +%Y-%m-%d_%H-%M-%S)"
  else
    parallel_base_dir="$(dirname "$json_filename")/out/parallel_$(date +%Y-%m-%d_%H-%M-%S)"
  fi
  mkdir -p "$parallel_base_dir"
  failure_sentinel="$parallel_base_dir/.failure"

  # Per-worker max_runs: split evenly (>=1 each). 0 means unlimited.
  if [ "$max_runs" -gt 0 ]; then
    worker_max_runs=$(( max_runs / parallel ))
    if [ "$worker_max_runs" -lt 1 ]; then
      worker_max_runs=1
    fi
  else
    worker_max_runs=0
  fi

  echo "Running $parallel parallel simulation workers ($worker_max_runs runs each, output: $parallel_base_dir)"

  # Build per-worker flags by filtering out --max_runs and --output-dir from
  # the existing args array (we set those per worker), and stripping off the
  # trailing positional JSON filename. Go's flag.Parse stops at the first
  # non-flag arg, so we must keep flags before the JSON filename in the
  # final invocation. Bash 3.2 compatible: indexed array with skip flag.
  worker_flags=()
  skip_next=0
  for a in "${args[@]}"; do
    if [ "$skip_next" -eq 1 ]; then
      skip_next=0
      continue
    fi
    case "$a" in
      --max_runs|--output-dir)
        skip_next=1
        ;;
      "$json_filename")
        # Skip the positional JSON filename; we append it after the per-
        # worker flags below.
        ;;
      *)
        worker_flags+=("$a")
        ;;
    esac
  done

  # Spawn N workers. Each runs the fizzbee binary on the (already-parsed)
  # json file, with per-worker --max_runs and --output-dir.
  pids=()
  i=0
  while [ "$i" -lt "$parallel" ]; do
    worker_out="$parallel_base_dir/worker_$i"
    worker_log="$parallel_base_dir/worker_$i.log"
    mkdir -p "$worker_out"
    (
      "$FIZZBEE_BIN" "${worker_flags[@]}" \
        --max_runs "$worker_max_runs" --output-dir "$worker_out" \
        "$json_filename" \
        > "$worker_log" 2>&1
      # Binary doesn't exit non-zero on FAILED — detect from output.
      if grep -qE "^FAILED|^DEADLOCK" "$worker_log" 2>/dev/null; then
        # First-writer-wins is fine; sentinel just signals "someone failed".
        touch "$failure_sentinel"
      fi
    ) &
    pids+=("$!")
    i=$(( i + 1 ))
  done

  # Ctrl-C: kill all live workers, exit 130 (SIGINT convention).
  trap '
    echo
    echo "Interrupted. Killing workers..."
    for p in "${pids[@]}"; do
      kill "$p" 2>/dev/null
    done
    exit 130
  ' INT TERM

  # Polling loop: wait for failure sentinel OR all workers to finish.
  while true; do
    if [ -f "$failure_sentinel" ]; then
      # Someone failed — kill survivors.
      for p in "${pids[@]}"; do
        kill "$p" 2>/dev/null
      done
      break
    fi
    alive=0
    for p in "${pids[@]}"; do
      if kill -0 "$p" 2>/dev/null; then
        alive=$(( alive + 1 ))
      fi
    done
    if [ "$alive" -eq 0 ]; then
      break
    fi
    sleep 0.5
  done

  # Reap all workers (including any killed by the survivors loop).
  for p in "${pids[@]}"; do
    wait "$p" 2>/dev/null
  done

  # Aggregate result.
  rm -f "$temp_output"
  if [ -f "$failure_sentinel" ]; then
    # Print first failing worker's full log (iterate in worker-id order).
    i=0
    while [ "$i" -lt "$parallel" ]; do
      log="$parallel_base_dir/worker_$i.log"
      if [ -f "$log" ] && grep -qE "^FAILED|^DEADLOCK" "$log" 2>/dev/null; then
        echo "=== Worker $i failed (output dir: $parallel_base_dir/worker_$i) ==="
        cat "$log"
        exit 1
      fi
      i=$(( i + 1 ))
    done
    # Sentinel set but no log matched (race / killed mid-write); generic msg.
    echo "FAILED: a worker reported failure but no log captured the trace. Check $parallel_base_dir"
    exit 1
  fi

  # All success — sum "Stopped after N runs" across workers.
  total=0
  i=0
  while [ "$i" -lt "$parallel" ]; do
    log="$parallel_base_dir/worker_$i.log"
    n=$(grep -oE "Stopped after [0-9]+" "$log" 2>/dev/null | grep -oE "[0-9]+" | head -1)
    if [ -n "$n" ]; then
      total=$(( total + n ))
    fi
    i=$(( i + 1 ))
  done
  echo "PASSED: $total simulation runs across $parallel workers"
  exit 0
fi
# --- END PARALLEL SIMULATION HANDLER ---------------------------------------


# Run the second command with the JSON filename
"$FIZZBEE_BIN" "${args[@]}"

# Clean up the temporary file
rm "$temp_output"