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

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

TEST_ROOT="./test_merged"
OUTPUT_DIR="./lab4_instcount_results"
COMPILER="./build/bin/compiler"
VERIFY_SCRIPT=""
MAX_CASES=0
STOP_ON_FIRST_FAILURE=false
START_FROM=1
KEEP_OLD=false
TEST_MODE="ir"
RUN_TESTS=false

show_help() {
  cat << 'EOF'
用法: ./test4_instcount.sh [选项]

说明:
  批量执行 Lab4(基本标量优化)测试。
  比较优化前后的指令数量变化，而非运行时间。
  这是一个更稳定的优化效果评估方法。

选项:
  -h, --help              显示此帮助信息
  -n, --max-cases N       最多运行 N 个测试用例 (0=不限制)
  -x, --stop-on-failure   遇到第一个失败即停止
  -s, --start-from N      从第 N 个测试开始
  -k, --keep              保留旧输出目录，不删除
  -t, --test-root DIR     指定测试目录
  -o, --output-dir DIR    指定输出目录
  -m, --mode MODE         测试模式: ir 或 asm (默认: ir)
                          ir: 统计 IR 指令数量
                          asm: 统计汇编指令数量
  -r, --run               运行测试用例验证正确性（默认不运行）

示例:
  ./test4_instcount.sh                    # 仅统计指令数量 (IR模式)
  ./test4_instcount.sh -r                 # 统计指令数量并运行验证
  ./test4_instcount.sh -m asm             # ASM模式，仅统计汇编指令
  ./test4_instcount.sh -r -m asm          # ASM模式，统计并运行
  ./test4_instcount.sh -n 10              # 只运行前10个
  ./test4_instcount.sh -s --mode ir       # IR模式，失败即停止
EOF
}

parse_args() {
  while [[ $# -gt 0 ]]; do
    case $1 in
      -h|--help)
        show_help
        exit 0
        ;;
      -n|--max-cases)
        MAX_CASES="$2"
        shift 2
        ;;
      -x|--stop-on-failure)
        STOP_ON_FIRST_FAILURE=true
        shift
        ;;
      -s|--start-from)
        START_FROM="$2"
        shift 2
        ;;
      -k|--keep)
        KEEP_OLD=true
        shift
        ;;
      -t|--test-root)
        TEST_ROOT="$2"
        shift 2
        ;;
      -o|--output-dir)
        OUTPUT_DIR="$2"
        shift 2
        ;;
      -m|--mode)
        TEST_MODE="$2"
        if [[ "$TEST_MODE" != "ir" && "$TEST_MODE" != "asm" ]]; then
          echo -e "${RED}错误: 模式必须是 'ir' 或 'asm'${NC}"
          exit 1
        fi
        shift 2
        ;;
      -r|--run)
        RUN_TESTS=true
        shift
        ;;
      *)
        echo -e "${RED}未知参数: $1${NC}"
        show_help
        exit 1
        ;;
    esac
  done
}

parse_args "$@"

if [[ "$TEST_MODE" == "ir" ]]; then
  VERIFY_SCRIPT="./scripts/verify_ir.sh"
else
  VERIFY_SCRIPT="./scripts/verify_asm.sh"
fi

if [[ ! -x "$COMPILER" ]]; then
  echo -e "${RED}错误: 编译器不存在或不可执行: $COMPILER${NC}"
  echo "请先编译项目: cd build && cmake .. && make -j4"
  exit 1
fi

if [[ "$RUN_TESTS" == "true" && ! -x "$VERIFY_SCRIPT" ]]; then
  echo -e "${RED}错误: verify脚本不可执行: $VERIFY_SCRIPT${NC}"
  exit 1
fi

if [[ ! -d "$TEST_ROOT" ]]; then
  echo -e "${RED}错误: 测试目录不存在: $TEST_ROOT${NC}"
  exit 1
fi

if [[ "$KEEP_OLD" != "true" ]]; then
  rm -rf "$OUTPUT_DIR"
fi
mkdir -p "$OUTPUT_DIR"

OUTPUT_DIR="$OUTPUT_DIR/$TEST_MODE"
mkdir -p "$OUTPUT_DIR"

LOG_FILE="$OUTPUT_DIR/lab4_instcount.log"
FAIL_FILE="$OUTPUT_DIR/failed_cases.txt"
ERROR_LOG_FILE="$OUTPUT_DIR/error_log.txt"
INST_REPORT_FILE="$OUTPUT_DIR/instruction_report.txt"
INST_SUMMARY_FILE="$OUTPUT_DIR/instruction_summary.txt"

if [[ "$KEEP_OLD" != "true" ]]; then
  : > "$LOG_FILE"
  : > "$FAIL_FILE"
  : > "$ERROR_LOG_FILE"
  : > "$INST_SUMMARY_FILE"
fi

{
  echo "================================================"
  echo "Lab4 指令数量对比报告"
  echo "================================================"
  echo "测试模式: ${TEST_MODE^^}"
  echo "运行验证: ${RUN_TESTS}"
  echo "生成时间: $(date)"
  echo "================================================"
  echo ""
  printf "%-60s | %12s | %12s | %10s\n" "测试用例" "原始指令数" "优化后指令数" "变化率"
  echo "------------------------------------------------"
} > "$INST_REPORT_FILE"

echo "Lab4 指令数量测试日志 - $(date)" >> "$LOG_FILE"
echo "TEST_ROOT=$TEST_ROOT" >> "$LOG_FILE"
echo "OUTPUT_DIR=$OUTPUT_DIR" >> "$LOG_FILE"
echo "TEST_MODE=${TEST_MODE^^}" >> "$LOG_FILE"
echo "RUN_TESTS=$RUN_TESTS" >> "$LOG_FILE"
echo "MAX_CASES=$MAX_CASES" >> "$LOG_FILE"
echo "================================================" >> "$LOG_FILE"

collect_sy_files() {
  find "$TEST_ROOT" -type f -name '*.sy' | sort
}

count_ir_instructions() {
  local ir_file="$1"
  if [[ ! -f "$ir_file" ]]; then
    echo "0"
    return
  fi
  
  local count=0
  local term_count=0
  local store_count=0
  local call_count=0
  
  count=$(grep -cE '^\s*%?[a-zA-Z0-9_]+\s*=' "$ir_file" 2>/dev/null || echo "0")
  count=$(echo "$count" | tr -d '[:space:]')
  [[ ! "$count" =~ ^[0-9]+$ ]] && count=0
  
  term_count=$(grep -cE '^\s*(ret|br|switch)' "$ir_file" 2>/dev/null || echo "0")
  term_count=$(echo "$term_count" | tr -d '[:space:]')
  [[ ! "$term_count" =~ ^[0-9]+$ ]] && term_count=0
  
  store_count=$(grep -cE '^\s*store ' "$ir_file" 2>/dev/null || echo "0")
  store_count=$(echo "$store_count" | tr -d '[:space:]')
  [[ ! "$store_count" =~ ^[0-9]+$ ]] && store_count=0
  
  call_count=$(grep -cE '^\s*(call|invoke)' "$ir_file" 2>/dev/null || echo "0")
  call_count=$(echo "$call_count" | tr -d '[:space:]')
  [[ ! "$call_count" =~ ^[0-9]+$ ]] && call_count=0
  
  echo $((count + term_count + store_count + call_count))
}

count_asm_instructions() {
  local asm_file="$1"
  if [[ ! -f "$asm_file" ]]; then
    echo "0"
    return
  fi
  
  local count=0
  count=$(grep -cE '^\s+[a-z]+' "$asm_file" 2>/dev/null || echo "0")
  count=$(echo "$count" | tr -d '[:space:]')
  [[ ! "$count" =~ ^[0-9]+$ ]] && count=0
  
  echo "$count"
}

run_single_test() {
  local sy_file="$1"
  local case_num="$2"
  local filename=$(basename "$sy_file")
  local base_name="${filename%.sy}"
  
  local rel_path="${sy_file#$TEST_ROOT/}"
  local rel_dir=$(dirname "$rel_path")
  
  if [[ "$rel_dir" == "." ]]; then
    rel_dir=""
  fi
  
  local case_out_dir="$OUTPUT_DIR"
  if [[ -n "$rel_dir" ]]; then
    case_out_dir="$OUTPUT_DIR/$rel_dir"
  fi
  mkdir -p "$case_out_dir"
  
  echo -ne "${YELLOW}[$case_num] $filename ... ${NC}"
  
  local orig_file=""
  local opt_file=""
  
  if [[ "$TEST_MODE" == "ir" ]]; then
    orig_file="$case_out_dir/${base_name}_orig.ll"
    opt_file="$case_out_dir/${base_name}_opt.ll"
    
    set +e
    "$COMPILER" --emit-ir "$sy_file" > "$orig_file" 2>/dev/null
    local orig_code=$?
    "$COMPILER" -O --emit-ir "$sy_file" > "$opt_file" 2>/dev/null
    local opt_code=$?
    set -e
    
    if [[ $orig_code -ne 0 || $opt_code -ne 0 ]]; then
      echo -e "${RED}编译失败${NC}"
      echo "[FAILED] $sy_file (compile error)" >> "$LOG_FILE"
      FAILED=$((FAILED + 1))
      echo "$sy_file" >> "$FAIL_FILE"
      if [[ "$STOP_ON_FIRST_FAILURE" = true ]]; then
        return 1
      fi
      return 0
    fi
    
    if [[ "$RUN_TESTS" == "true" ]]; then
      "$VERIFY_SCRIPT" "$sy_file" "$case_out_dir" -O --run > /dev/null 2>&1
      local run_code=$?
      
      if [[ $run_code -ne 0 ]]; then
        echo -e "${RED}运行失败${NC}"
        echo "[FAILED] $sy_file (runtime error, exit=$run_code)" >> "$LOG_FILE"
        FAILED=$((FAILED + 1))
        echo "$sy_file" >> "$FAIL_FILE"
        if [[ "$STOP_ON_FIRST_FAILURE" = true ]]; then
          return 1
        fi
        return 0
      fi
    fi
  else
    orig_file="$case_out_dir/${base_name}_orig.s"
    opt_file="$case_out_dir/${base_name}_opt.s"
    
    set +e
    "$COMPILER" --emit-asm "$sy_file" > "$orig_file" 2>/dev/null
    local orig_code=$?
    "$COMPILER" -O --emit-asm "$sy_file" > "$opt_file" 2>/dev/null
    local opt_code=$?
    set -e
    
    if [[ $orig_code -ne 0 || $opt_code -ne 0 ]]; then
      echo -e "${RED}编译失败${NC}"
      echo "[FAILED] $sy_file (compile error)" >> "$LOG_FILE"
      FAILED=$((FAILED + 1))
      echo "$sy_file" >> "$FAIL_FILE"
      if [[ "$STOP_ON_FIRST_FAILURE" = true ]]; then
        return 1
      fi
      return 0
    fi
    
    if [[ "$RUN_TESTS" == "true" ]]; then
      "$VERIFY_SCRIPT" "$sy_file" "$case_out_dir" -O --run > /dev/null 2>&1
      local run_code=$?
      
      if [[ $run_code -ne 0 ]]; then
        echo -e "${RED}运行失败${NC}"
        echo "[FAILED] $sy_file (runtime error, exit=$run_code)" >> "$LOG_FILE"
        FAILED=$((FAILED + 1))
        echo "$sy_file" >> "$FAIL_FILE"
        if [[ "$STOP_ON_FIRST_FAILURE" = true ]]; then
          return 1
        fi
        return 0
      fi
    fi
  fi
  
  local orig_count=0
  local opt_count=0
  
  if [[ "$TEST_MODE" == "ir" ]]; then
    orig_count=$(count_ir_instructions "$orig_file")
    opt_count=$(count_ir_instructions "$opt_file")
  else
    orig_count=$(count_asm_instructions "$orig_file")
    opt_count=$(count_asm_instructions "$opt_file")
  fi
  
  local change_rate=""
  if [[ "$orig_count" -gt 0 ]]; then
    change_rate=$(awk -v orig="$orig_count" -v opt="$opt_count" 'BEGIN { printf "%.2f", ((orig - opt) / orig) * 100 }')
  fi
  
  SUCCESS=$((SUCCESS + 1))
  
  total_orig_inst=$((total_orig_inst + orig_count))
  total_opt_inst=$((total_opt_inst + opt_count))
  inst_cases_count=$((inst_cases_count + 1))
  
  if [[ -n "$change_rate" ]]; then
    if (( $(echo "$change_rate > 0" | bc -l 2>/dev/null || echo "0") )); then
      echo -ne "${GREEN}成功${NC} "
      echo -ne "${GREEN}(指令: ${orig_count} -> ${opt_count}, -${change_rate}%)${NC}"
      total_reduction_sum=$(echo "$total_reduction_sum + $change_rate" | bc -l 2>/dev/null || echo "0")
      reduction_count=$((reduction_count + 1))
    elif (( $(echo "$change_rate < 0" | bc -l 2>/dev/null || echo "0") )); then
      echo -ne "${GREEN}成功${NC} "
      echo -ne "${RED}(指令: ${orig_count} -> ${opt_count}, +${change_rate#-}%)${NC}"
      total_increase_sum=$(echo "$total_increase_sum + (0 - $change_rate)" | bc -l 2>/dev/null || echo "0")
      increase_count=$((increase_count + 1))
    else
      echo -ne "${GREEN}成功${NC} "
      echo -ne "(指令: ${orig_count} -> ${opt_count}, 0.00%)${NC}"
    fi
    
    printf "%-60s | %12d | %12d | %9s%%\n" "$rel_path" "$orig_count" "$opt_count" "$change_rate" >> "$INST_REPORT_FILE"
    echo "$rel_path: orig=$orig_count, opt=$opt_count, change=${change_rate}%" >> "$INST_SUMMARY_FILE"
  else
    echo -e "${GREEN}成功${NC} (指令: ${orig_count} -> ${opt_count})"
    printf "%-60s | %12d | %12d | %10s\n" "$rel_path" "$orig_count" "$opt_count" "N/A" >> "$INST_REPORT_FILE"
  fi
  echo ""
  
  echo "[SUCCESS] $sy_file (orig=$orig_count, opt=$opt_count)" >> "$LOG_FILE"
  
  return 0
}

main() {
  local sy_files=($(collect_sy_files))
  local total=${#sy_files[@]}
  
  if [[ $total -eq 0 ]]; then
    echo -e "${YELLOW}警告: 未找到任何测试文件${NC}"
    exit 0
  fi
  
  echo -e "${BLUE}========================================${NC}"
  echo -e "${BLUE}Lab4 指令数量对比测试 (${TEST_MODE^^}模式)${NC}"
  echo -e "${BLUE}========================================${NC}"
  echo "测试目录: $TEST_ROOT"
  echo "输出目录: $OUTPUT_DIR"
  echo "运行验证: $RUN_TESTS"
  echo "测试用例总数: $total"
  echo ""
  
  SUCCESS=0
  FAILED=0
  case_num=0
  
  total_orig_inst=0
  total_opt_inst=0
  inst_cases_count=0
  total_reduction_sum=0
  reduction_count=0
  total_increase_sum=0
  increase_count=0
  
  for sy_file in "${sy_files[@]}"; do
    case_num=$((case_num + 1))
    
    if [[ $case_num -lt $START_FROM ]]; then
      continue
    fi
    
    local executed=$((case_num - START_FROM + 1))
    if [[ $MAX_CASES -gt 0 && $executed -gt $MAX_CASES ]]; then
      break
    fi
    
    set +e
    run_single_test "$sy_file" "$case_num"
    test_result=$?
    set -e
    
    if [[ $test_result -eq 1 ]] && [[ "$STOP_ON_FIRST_FAILURE" = true ]]; then
      break
    fi
  done
  
  {
    echo ""
    echo "================================================"
    echo "指令数量对比摘要"
    echo "================================================"
    echo "总原始指令数: $total_orig_inst"
    echo "总优化后指令数: $total_opt_inst"
    
    if [[ $total_orig_inst -gt 0 ]]; then
      local total_change=$(awk -v orig="$total_orig_inst" -v opt="$total_opt_inst" 'BEGIN { printf "%.2f", ((orig - opt) / orig) * 100 }')
      echo "总体指令变化率: ${total_change}%"
    fi
    
    echo ""
    echo "测试用例统计:"
    echo "  指令减少的用例: $reduction_count"
    echo "  指令增加的用例: $increase_count"
    echo "  指令不变的用例: $((inst_cases_count - reduction_count - increase_count))"
    
    if [[ $reduction_count -gt 0 ]]; then
      local avg_reduction=$(echo "scale=2; $total_reduction_sum / $reduction_count" | bc 2>/dev/null || echo "0")
      echo "  平均指令减少率: ${avg_reduction}%"
    fi
    
    if [[ $increase_count -gt 0 ]]; then
      local avg_increase=$(echo "scale=2; $total_increase_sum / $increase_count" | bc 2>/dev/null || echo "0")
      echo "  平均指令增加率: ${avg_increase}%"
    fi
    
    echo "================================================"
  } >> "$INST_REPORT_FILE"
  
  RATE="0.00"
  if [[ $case_num -gt 0 ]]; then
    RATE=$(awk -v s="$SUCCESS" -v t="$case_num" 'BEGIN { printf "%.2f", (s*100.0)/t }')
  fi
  
  echo ""
  echo -e "${BLUE}========================================${NC}"
  echo -e "${BLUE}Lab4 指令数量测试完成${NC}"
  echo -e "${BLUE}执行用例: $case_num${NC}"
  echo -e "${GREEN}成功: $SUCCESS${NC}"
  echo -e "${RED}失败: $FAILED${NC}"
  echo -e "${BLUE}成功率: ${RATE}%${NC}"
  
  echo ""
  echo -e "${BLUE}指令数量统计:${NC}"
  echo "  总原始指令数: $total_orig_inst"
  echo "  总优化后指令数: $total_opt_inst"
  
  if [[ $total_orig_inst -gt 0 ]]; then
    local total_change=$(awk -v orig="$total_orig_inst" -v opt="$total_opt_inst" 'BEGIN { printf "%.2f", ((orig - opt) / orig) * 100 }')
    if (( $(echo "$total_change >= 0" | bc -l 2>/dev/null || echo "0") )); then
      echo -e "  总体变化率: ${GREEN}${total_change}%${NC}"
    else
      echo -e "  总体变化率: ${RED}${total_change}%${NC}"
    fi
  fi
  
  echo ""
  echo -e "${BLUE}用例统计:${NC}"
  echo -e "  ${GREEN}指令减少的用例: $reduction_count${NC}"
  echo -e "  ${RED}指令增加的用例: $increase_count${NC}"
  echo "  指令不变的用例: $((inst_cases_count - reduction_count - increase_count))"
  
  if [[ $reduction_count -gt 0 ]]; then
    local avg_reduction=$(echo "scale=2; $total_reduction_sum / $reduction_count" | bc 2>/dev/null || echo "0")
    echo -e "  ${GREEN}平均指令减少率: ${avg_reduction}%${NC}"
  fi
  
  if [[ $increase_count -gt 0 ]]; then
    local avg_increase=$(echo "scale=2; $total_increase_sum / $increase_count" | bc 2>/dev/null || echo "0")
    echo -e "  ${RED}平均指令增加率: ${avg_increase}%${NC}"
  fi
  
  echo ""
  echo -e "${BLUE}详细报告: $INST_REPORT_FILE${NC}"
  echo -e "${BLUE}日志: $LOG_FILE${NC}"
  if [[ $FAILED -gt 0 ]]; then
    echo -e "${RED}失败清单: $FAIL_FILE${NC}"
    echo -e "${RED}错误日志: $ERROR_LOG_FILE${NC}"
  fi
  echo -e "${BLUE}========================================${NC}"
  
  exit 0
}

main