#!/bin/bash
set -u
set -o pipefail

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
cd "$SCRIPT_DIR"

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m'

COMPILER=./build/bin/compiler
TESTS=(
  "huffman-01" "huffman-02" "huffman-03"
  "many_mat_cal-1" "many_mat_cal-2" "many_mat_cal-3"
  "h-1-01" "h-1-02" "h-1-03"
  "fft0" "fft1" "fft2"
  "matmul1" "matmul2" "matmul3"
  "conv2d-1" "conv2d-2" "conv2d-3"
  "shuffle0" "shuffle1" "shuffle2"
  "sl1" "sl2" "sl3"
  "crypto-1" "crypto-2" "crypto-3"
  "crc1" "crc2" "crc3"
  "h-4-01" "h-4-02" "h-4-03"
  "h-5-01" "h-5-02" "h-5-03"
  "h-8-01" "h-8-02" "h-8-03"
  "h-9-01" "h-9-02" "h-9-03"
  "h-10-01" "h-10-02" "h-10-03"
  "knapsack_naive-1" "knapsack_naive-2" "knapsack_naive-3"
  "optimization_scheduling1" "optimization_scheduling2" "optimization_scheduling3"
  "transpose0" "transpose1" "transpose2"
  "03_sort1" "03_sort2" "03_sort3"
  "01_mm1" "01_mm2" "01_mm3"
)
BASELINE_FILE="./指令数基线.md"
RESULT_DIR="./.count_results"
JOBS=$(nproc 2>/dev/null || echo 4)

# ============================================================
# Worker: 编译单个测试并计数
# ============================================================
run_one_count() {
  local idx="$1" name="$2"
  local result_file="$RESULT_DIR/$idx"
  local tmp_asm="$RESULT_DIR/.tmp_${idx}.s"

  set +e
  timeout --signal=KILL 60 "$COMPILER" -S -O -o "$tmp_asm" "2026test/performance/${name}.sy" 2>/dev/null
  local lines
  lines=$(wc -l < "$tmp_asm" 2>/dev/null)
  lines=${lines:-0}
  rm -f "$tmp_asm"
  set -e

  printf 'TEST=%s\nLINES=%s\n' "$name" "$lines" > "$result_file"
}

# ============================================================
# 构建带索引的测试队列
# ============================================================
declare -a TEST_QUEUE=()
for i in "${!TESTS[@]}"; do
  TEST_QUEUE+=("$i|${TESTS[$i]}")
done

mkdir -p "$RESULT_DIR"
wall_start=$(date +%s%3N 2>/dev/null || date +%s000)

# ============================================================
# 并行执行编译（或串行 fallback）
# ============================================================
if [[ $JOBS -gt 1 && ${#TEST_QUEUE[@]} -gt 1 ]]; then
  export RESULT_DIR COMPILER
  export -f run_one_count

  printf '%s\n' "${TEST_QUEUE[@]}" | xargs -P "$JOBS" -L 1 bash -c '
    IFS="|" read -r idx name <<< "$1"
    run_one_count "$idx" "$name"
  ' _
else
  for item in "${TEST_QUEUE[@]}"; do
    IFS='|' read -r idx name <<< "$item"
    run_one_count "$idx" "$name"
  done
fi

wall_end=$(date +%s%3N 2>/dev/null || date +%s000)
wall_elapsed=$((wall_end - wall_start))

# ============================================================
# 汇总结果（串行处理基线对比和文件更新）
# ============================================================
UPDATED=0
TOTAL_TESTS=0

echo "=== 指令数测试 ==="

for item in "${TEST_QUEUE[@]}"; do
  IFS='|' read -r idx name <<< "$item"
  result_file="$RESULT_DIR/$idx"

  if [[ ! -f "$result_file" ]]; then
    echo -e "${RED}performance/${name}: 编译失败/超时${NC}"
    continue
  fi

  test_name=$(grep '^TEST=' "$result_file" | cut -d= -f2)
  lines=$(grep '^LINES=' "$result_file" | cut -d= -f2)

  if [[ ! "$lines" =~ ^[0-9]+$ ]]; then
    echo -e "${RED}performance/${name}: 指令数异常 (${lines})${NC}"
    continue
  fi

  TOTAL_TESTS=$((TOTAL_TESTS + 1))

  baseline=$(grep "performance/${name}" "$BASELINE_FILE" 2>/dev/null | grep -oP '\|\s*\d+\s*\|' | grep -oP '\d+' | head -1)

  if [[ -n "$baseline" && "$baseline" =~ ^[0-9]+$ ]]; then
    if [[ "$lines" -lt "$baseline" ]]; then
      diff=$((baseline - lines))
      echo -e "performance/${name}: ${lines} ${GREEN}(基线:${baseline} 减少${diff}行) <- 更新基线${NC}"
      sed -i "s/| performance\/${name} | ${baseline} |/| performance\/${name} | ${lines} |/" "$BASELINE_FILE"
      UPDATED=$((UPDATED + 1))
    elif [[ "$lines" -gt "$baseline" ]]; then
      echo -e "performance/${name}: ${lines} ${RED}(基线:${baseline} 增加$((lines - baseline))行)${NC}"
    else
      echo "performance/${name}: ${lines} (基线:${baseline} 持平)"
    fi
  else
    echo "performance/${name}: ${lines}"
  fi
done

# ============================================================
# 清理
# ============================================================
rm -rf "$RESULT_DIR"

echo ""
echo -e "${BLUE}总耗时: ${wall_elapsed}ms | 测试: ${TOTAL_TESTS}/${#TESTS[@]}${NC}"

if [[ $UPDATED -gt 0 ]]; then
  echo -e "${GREEN}已更新 ${UPDATED} 个测试集的指令数基线${NC}"
fi