#!/usr/bin/env bash

set -euo pipefail

ROOT_DIR="$(cd "$(dirname "$0")/.." && pwd)"
BUILD_DIR="$ROOT_DIR/build"
ANTLR_DIR="$BUILD_DIR/generated/antlr4"
JAR_PATH="$ROOT_DIR/third_party/antlr-4.13.2-complete.jar"
GRAMMAR_PATH="$ROOT_DIR/src/antlr4/SysY.g4"

RUN_FUNCTIONAL=true
RUN_PERFORMANCE=true
RUN_IR=true
RUN_ASM=true
RUN_EXEC=true
DO_BUILD=true
RUN_TIMEOUT=""
OPT_LEVEL=0
OUT_ROOT=""
TIME_LOG_FILE=""
failed_cases=()

declare -A counters=()

usage() {
  cat <<'EOF'
Usage: ./solution/run_lab4_batch.sh [options]

Options:
  --no-build          Skip ANTLR generation and project rebuild
  --functional-only   Run only test/test_case/functional/*.sy
  --performance-only  Run only test/test_case/performance/*.sy
  --ir-only           Run only verify_ir batch
  --asm-only          Run only verify_asm batch
  --no-run            Generate IR/asm only; skip execution and output check
  --emit-only         Alias of --no-run
  --timeout <sec>     Apply per-case timeout via the `timeout` command
  -O0                 Run batch with compiler flag -O0 (default)
  -O1                 Run batch with compiler flag -O1
  --opt-level <0|1>   Same as -O0 / -O1
  --output-dir <dir>  Set output root; default is test/test_result/lab4_batch_o<level>
  --help              Show this help message
EOF
}

set_opt_level() {
  case "$1" in
    0|O0|-O0)
      OPT_LEVEL=0
      ;;
    1|O1|-O1)
      OPT_LEVEL=1
      ;;
    *)
      echo "Unsupported opt level: $1" >&2
      echo "Only -O0 / -O1 are supported." >&2
      exit 1
      ;;
  esac
}

bump_counter() {
  local mode=$1
  local group=$2
  local kind=$3
  local key="$mode:$group:$kind"
  counters["$key"]=$(( ${counters["$key"]:-0} + 1 ))
}

get_counter() {
  local mode=$1
  local group=$2
  local kind=$3
  echo "${counters["$mode:$group:$kind"]:-0}"
}

print_mode_summary() {
  local mode=$1
  local functional_total functional_passed functional_failed
  local performance_total performance_passed performance_failed
  local total passed failed

  functional_total=$(get_counter "$mode" "functional" "total")
  functional_passed=$(get_counter "$mode" "functional" "passed")
  functional_failed=$(get_counter "$mode" "functional" "failed")
  performance_total=$(get_counter "$mode" "performance" "total")
  performance_passed=$(get_counter "$mode" "performance" "passed")
  performance_failed=$(get_counter "$mode" "performance" "failed")

  total=$((functional_total + performance_total))
  passed=$((functional_passed + performance_passed))
  failed=$((functional_failed + performance_failed))

  echo "  ${mode^^} functional: total=$functional_total, passed=$functional_passed, failed=$functional_failed"
  echo "  ${mode^^} performance: total=$performance_total, passed=$performance_passed, failed=$performance_failed"
  echo "  ${mode^^} overall: total=$total, passed=$passed, failed=$failed"
}

print_summary() {
  local overall_total=0
  local overall_passed=0
  local overall_failed=0

  echo
  echo "Summary:"
  echo "  Opt level: -O$OPT_LEVEL"
  echo "  Execution: $([[ "$RUN_EXEC" == true ]] && echo "enabled" || echo "disabled")"

  if [[ "$RUN_IR" == true ]]; then
    print_mode_summary "ir"
    overall_total=$((overall_total + $(get_counter "ir" "functional" "total") + $(get_counter "ir" "performance" "total")))
    overall_passed=$((overall_passed + $(get_counter "ir" "functional" "passed") + $(get_counter "ir" "performance" "passed")))
    overall_failed=$((overall_failed + $(get_counter "ir" "functional" "failed") + $(get_counter "ir" "performance" "failed")))
  fi

  if [[ "$RUN_ASM" == true ]]; then
    print_mode_summary "asm"
    overall_total=$((overall_total + $(get_counter "asm" "functional" "total") + $(get_counter "asm" "performance" "total")))
    overall_passed=$((overall_passed + $(get_counter "asm" "functional" "passed") + $(get_counter "asm" "performance" "passed")))
    overall_failed=$((overall_failed + $(get_counter "asm" "functional" "failed") + $(get_counter "asm" "performance" "failed")))
  fi

  echo "  Overall: total=$overall_total, passed=$overall_passed, failed=$overall_failed"

  if (( ${#failed_cases[@]} > 0 )); then
    echo "Failed cases:"
    printf '  - %s\n' "${failed_cases[@]}"
  fi
  echo "  Per-case timing log: $TIME_LOG_FILE"
}

format_elapsed_seconds() {
  local elapsed_ns=$1
  awk -v ns="$elapsed_ns" 'BEGIN { printf "%.3f", ns / 1000000000 }'
}

run_case() {
  local mode=$1
  local group=$2
  local case_file=$3
  local stem out_dir log_file
  local start_ns end_ns elapsed_ns elapsed_s
  local -a cmd

  stem="$(basename "${case_file%.sy}")"
  out_dir="$OUT_ROOT/$mode/$group"
  log_file="$out_dir/$stem.verify.log"
  mkdir -p "$out_dir"

  bump_counter "$mode" "$group" "total"

  if [[ "$mode" == "ir" ]]; then
    cmd=("$ROOT_DIR/scripts/verify_ir.sh" "$case_file" "$out_dir")
  else
    cmd=("$ROOT_DIR/scripts/verify_asm.sh" "$case_file" "$out_dir")
  fi

  if [[ "$RUN_EXEC" == true ]]; then
    cmd+=(--run)
  fi
  cmd+=(-- "-O$OPT_LEVEL")

  if [[ -n "$RUN_TIMEOUT" ]]; then
    cmd=(timeout "$RUN_TIMEOUT" "${cmd[@]}")
  fi

  start_ns=$(date +%s%N)
  if "${cmd[@]}" >"$log_file" 2>&1; then
    end_ns=$(date +%s%N)
    elapsed_ns=$((end_ns - start_ns))
    elapsed_s=$(format_elapsed_seconds "$elapsed_ns")
    echo "PASS [$mode] $case_file (${elapsed_s}s)"
    bump_counter "$mode" "$group" "passed"
    printf '%s,%s,%s,%s,%s,%s,%s\n' \
      "$mode" "$group" "$case_file" "PASS" "$elapsed_ns" "$elapsed_s" "$log_file" >> "$TIME_LOG_FILE"
  else
    end_ns=$(date +%s%N)
    elapsed_ns=$((end_ns - start_ns))
    elapsed_s=$(format_elapsed_seconds "$elapsed_ns")
    echo "FAIL [$mode] $case_file (${elapsed_s}s)"
    cat "$log_file"
    bump_counter "$mode" "$group" "failed"
    failed_cases+=("[$mode] $case_file")
    printf '%s,%s,%s,%s,%s,%s,%s\n' \
      "$mode" "$group" "$case_file" "FAIL" "$elapsed_ns" "$elapsed_s" "$log_file" >> "$TIME_LOG_FILE"
  fi
}

run_group() {
  local mode=$1
  local group=$2
  local case_dir=$3

  while IFS= read -r case_file; do
    run_case "$mode" "$group" "$case_file"
  done < <(find "$case_dir" -maxdepth 1 -type f -name '*.sy' | sort)
}

while [[ $# -gt 0 ]]; do
  case "$1" in
    --no-build)
      DO_BUILD=false
      ;;
    --functional-only)
      RUN_FUNCTIONAL=true
      RUN_PERFORMANCE=false
      ;;
    --performance-only)
      RUN_FUNCTIONAL=false
      RUN_PERFORMANCE=true
      ;;
    --ir-only)
      RUN_IR=true
      RUN_ASM=false
      ;;
    --asm-only)
      RUN_IR=false
      RUN_ASM=true
      ;;
    --no-run|--emit-only)
      RUN_EXEC=false
      ;;
    --timeout)
      shift
      if [[ $# -eq 0 ]]; then
        echo "Missing value for --timeout" >&2
        usage
        exit 1
      fi
      RUN_TIMEOUT="$1"
      ;;
    -O0|-O1)
      set_opt_level "$1"
      ;;
    --opt-level)
      shift
      if [[ $# -eq 0 ]]; then
        echo "Missing value for --opt-level" >&2
        usage
        exit 1
      fi
      set_opt_level "$1"
      ;;
    --output-dir)
      shift
      if [[ $# -eq 0 ]]; then
        echo "Missing value for --output-dir" >&2
        usage
        exit 1
      fi
      if [[ "$1" = /* ]]; then
        OUT_ROOT="$1"
      else
        OUT_ROOT="$ROOT_DIR/$1"
      fi
      ;;
    --help)
      usage
      exit 0
      ;;
    *)
      echo "Unknown option: $1" >&2
      usage
      exit 1
      ;;
  esac
  shift
done

if [[ -z "$OUT_ROOT" ]]; then
  OUT_ROOT="$ROOT_DIR/test/test_result/lab4_batch_o$OPT_LEVEL"
fi
mkdir -p "$OUT_ROOT"
TIME_LOG_FILE="$OUT_ROOT/case_timing.csv"
echo "mode,group,case,status,elapsed_ns,elapsed_s,log_file" > "$TIME_LOG_FILE"

if [[ "$RUN_FUNCTIONAL" == false && "$RUN_PERFORMANCE" == false ]]; then
  echo "No test set selected." >&2
  exit 1
fi

if [[ "$RUN_IR" == false && "$RUN_ASM" == false ]]; then
  echo "No verification pipeline selected." >&2
  exit 1
fi

if [[ -n "$RUN_TIMEOUT" ]] && ! command -v timeout >/dev/null 2>&1; then
  echo "未找到 timeout 命令，无法使用 --timeout。" >&2
  exit 1
fi

cd "$ROOT_DIR"

if [[ "$DO_BUILD" == true ]]; then
  echo "[1/4] Generating ANTLR sources..."
  mkdir -p "$ANTLR_DIR"
  java -jar "$JAR_PATH" \
    -Dlanguage=Cpp \
    -visitor -no-listener \
    -Xexact-output-dir \
    -o "$ANTLR_DIR" \
    "$GRAMMAR_PATH"

  echo "[2/4] Configuring CMake..."
  cmake -S "$ROOT_DIR" -B "$BUILD_DIR" -DCMAKE_BUILD_TYPE=Release -DCOMPILER_PARSE_ONLY=OFF

  echo "[3/4] Building project..."
  cmake --build "$BUILD_DIR" -j "$(nproc)"
fi

echo "[4/4] Running Lab4 batch tests..."

if [[ "$RUN_IR" == true ]]; then
  if [[ "$RUN_FUNCTIONAL" == true ]]; then
    run_group "ir" "functional" "$ROOT_DIR/test/test_case/functional"
  fi
  if [[ "$RUN_PERFORMANCE" == true ]]; then
    run_group "ir" "performance" "$ROOT_DIR/test/test_case/performance"
  fi
fi

if [[ "$RUN_ASM" == true ]]; then
  if [[ "$RUN_FUNCTIONAL" == true ]]; then
    run_group "asm" "functional" "$ROOT_DIR/test/test_case/functional"
  fi
  if [[ "$RUN_PERFORMANCE" == true ]]; then
    run_group "asm" "performance" "$ROOT_DIR/test/test_case/performance"
  fi
fi

print_summary

if (( ${#failed_cases[@]} > 0 )); then
  echo "Batch test finished with failures."
  exit 1
fi

echo "Batch test passed."