#!/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="./lab5_results"
COMPILER="./build/bin/compiler"
VERIFY_SCRIPT="./scripts/verify_asm.sh"
BASELINE_DIR="./lab3_results"
MAX_CASES=0
STOP_ON_FIRST_FAILURE=false
START_FROM=1
KEEP_OLD=false

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

说明:
  批量执行 Lab5(寄存器分配优化)测试。
  默认会递归扫描测试目录下所有 .sy 文件，
  使用寄存器分配优化编译并运行，验证输出正确性。
  结果保存在 lab5_results，与 lab3_results 对比。

选项:
  -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    指定输出目录

示例:
  ./test5.sh                          # 运行所有测试
  ./test5.sh -n 10                    # 只运行前10个
  ./test5.sh -s                       # 失败即停止
  ./test5.sh --start-from 5           # 从第5个开始
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
        ;;
      *)
        echo -e "${RED}未知参数: $1${NC}"
        show_help
        exit 1
        ;;
    esac
  done
}

parse_args "$@"

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

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

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

# 创建输出目录
if [[ "$KEEP_OLD" != "true" ]]; then
  rm -rf "$OUTPUT_DIR"
fi
mkdir -p "$OUTPUT_DIR"

# 初始化统计文件
LOG_FILE="$OUTPUT_DIR/lab5_batch.log"
FAIL_FILE="$OUTPUT_DIR/failed_cases.txt"
ERROR_LOG_FILE="$OUTPUT_DIR/error_log.txt"
TIME_SUMMARY_FILE="$OUTPUT_DIR/time_summary.txt"
OPT_REPORT_FILE="$OUTPUT_DIR/optimization_report.txt"

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

# 初始化优化报告
{
  echo "================================================"
  echo "Lab5 寄存器分配优化对比报告"
  echo "================================================"
  echo "测试模式: ASM"
  echo "基线目录: $BASELINE_DIR"
  echo "生成时间: $(date)"
  echo "================================================"
  echo ""
  echo "格式: 测试用例 | 基线时间(ms) | 优化后时间(ms) | 提升幅度"
  echo "------------------------------------------------"
} > "$OPT_REPORT_FILE"

echo "Lab5 批量测试日志 - $(date)" >> "$LOG_FILE"
echo "TEST_ROOT=$TEST_ROOT" >> "$LOG_FILE"
echo "OUTPUT_DIR=$OUTPUT_DIR" >> "$LOG_FILE"
echo "BASELINE_DIR=$BASELINE_DIR" >> "$LOG_FILE"
echo "MAX_CASES=$MAX_CASES" >> "$LOG_FILE"
echo "================================================" >> "$LOG_FILE"

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

# 计算优化提升百分比
calc_improvement() {
  local old_time="$1"
  local new_time="$2"
  
  if [[ -z "$old_time" || -z "$new_time" ]]; then
    echo ""
    return
  fi
  
  if [[ "$old_time" -eq 0 ]]; then
    echo ""
    return
  fi
  
  local improvement=$(awk -v old="$old_time" -v new="$new_time" 'BEGIN { printf "%.2f", ((old - new) / old) * 100 }')
  echo "$improvement"
}

run_single_test() {
  local sy_file="$1"
  local case_num="$2"
  local filename=$(basename "$sy_file")
  local base_name="${filename%.sy}"
  
  # 获取相对路径（去掉 TEST_ROOT 前缀）
  local rel_path="${sy_file#$TEST_ROOT/}"
  local rel_dir=$(dirname "$rel_path")
  
  # 如果 rel_dir 是 "."，则设为空字符串
  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"
  
  # 删除旧的时间文件
  rm -f "$case_out_dir/${base_name}_time.txt"
  
  echo -ne "${YELLOW}[$case_num] $filename ... ${NC}"
  
  # 调用 verify 脚本（临时禁用 set -e）
  set +e
  "$VERIFY_SCRIPT" "$sy_file" "$case_out_dir" --run > /dev/null 2>&1
  code=$?
  set -e
  
  # 读取优化后的运行时间
  local opt_time_file="$case_out_dir/${base_name}_time.txt"
  local opt_time=""
  if [[ -f "$opt_time_file" ]]; then
    opt_time=$(cat "$opt_time_file")
  fi
  
  # 读取基线时间
  local baseline_time=""
  local baseline_time_file=""
  
  # 构建基线时间文件的搜索路径
  local baseline_search_paths=(
    "$BASELINE_DIR/$rel_path/${base_name}_time.txt"
    "$BASELINE_DIR/$rel_dir/${base_name}_time.txt"
  )
  # 如果上面找不到，则用 find 查找
  if [[ ! -f "${baseline_search_paths[0]}" ]] && [[ ! -f "${baseline_search_paths[1]}" ]]; then
    baseline_time_file=$(find "$BASELINE_DIR" -name "${base_name}_time.txt" 2>/dev/null | head -n 1)
  fi
  
  # 尝试从搜索路径中查找
  for path in "${baseline_search_paths[@]}"; do
    if [[ -f "$path" ]]; then
      baseline_time_file="$path"
      break
    fi
  done
  
  if [[ -f "$baseline_time_file" ]]; then
    baseline_time=$(cat "$baseline_time_file")
  fi
  
  if [[ $code -eq 0 ]]; then
    SUCCESS=$((SUCCESS + 1))
    
    # 显示运行时间和优化对比
    if [[ -n "$opt_time" && "$opt_time" =~ ^[0-9]+$ ]]; then
      # 累加统计
      total_time_sum=$((total_time_sum + opt_time))
      time_cases_count=$((time_cases_count + 1))
      
      # 写入汇总文件
      echo "$rel_path: ${opt_time}ms" >> "$TIME_SUMMARY_FILE"
      
      # 先打印"成功"，再打印时间对比
      echo -ne "${GREEN}成功${NC} "
      
      # 优化对比
      if [[ -n "$baseline_time" && "$baseline_time" =~ ^[0-9]+$ ]]; then
        improvement=$(calc_improvement "$baseline_time" "$opt_time")
        if [[ -n "$improvement" ]]; then
          if (( $(echo "$improvement > 0" | bc -l 2>/dev/null || echo "0") )); then
            echo -ne "${GREEN}(${baseline_time}ms->${opt_time}ms，+${improvement}%)${NC}"
            total_improvement_sum=$(echo "$total_improvement_sum + $improvement" | bc -l 2>/dev/null || echo "0")
            improvement_count=$((improvement_count + 1))
          elif (( $(echo "$improvement < 0" | bc -l 2>/dev/null || echo "0") )); then
            echo -ne "${RED}(${baseline_time}ms->${opt_time}ms，${improvement}%)${NC}"
            total_degradation_sum=$(echo "$total_degradation_sum + (0 - $improvement)" | bc -l 2>/dev/null || echo "0")
            degradation_count=$((degradation_count + 1))
          else
            echo -ne "(${baseline_time}ms->${opt_time}ms，+0.00%)"
          fi
          
          # 写入优化报告
          echo "$rel_path | $baseline_time | $opt_time | ${improvement}%" >> "$OPT_REPORT_FILE"
        else
          echo -ne "(${baseline_time}ms->${opt_time}ms)"
        fi
      else
        echo -ne "(${opt_time}ms，无基线数据)"
      fi
      echo ""
    else
      echo -e "${GREEN}成功${NC} (运行时间: N/A)"
    fi
    
    echo "[SUCCESS] $sy_file" >> "$LOG_FILE"
  else
    FAILED=$((FAILED + 1))
    echo "$sy_file" >> "$FAIL_FILE"
    
    # 先打印"失败"，再显示时间（如果有）
    if [[ -n "$opt_time" && "$opt_time" =~ ^[0-9]+$ ]]; then
      echo -e "${RED}失败${NC} (${opt_time}ms)"
    else
      echo -e "${RED}失败${NC}"
    fi
    
    echo "[FAILED] $sy_file (exit=$code)" >> "$LOG_FILE"
    
    # 写入错误日志
    echo "========================================" >> "$ERROR_LOG_FILE"
    echo "测试失败: $sy_file" >> "$ERROR_LOG_FILE"
    echo "退出码: $code" >> "$ERROR_LOG_FILE"
    echo "时间: $(date)" >> "$ERROR_LOG_FILE"
    echo "========================================" >> "$ERROR_LOG_FILE"
    
    if [[ "$STOP_ON_FIRST_FAILURE" = true ]]; then
      echo -e "${RED}========================================${NC}"
      echo -e "${RED}在第一个失败处停止测试${NC}"
      echo -e "${RED}失败文件: $sy_file${NC}"
      echo -e "${RED}日志: $LOG_FILE${NC}"
      echo -e "${RED}错误日志: $ERROR_LOG_FILE${NC}"
      echo -e "${RED}========================================${NC}"
      return 1
    fi
  fi
  
  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}Lab5 寄存器分配优化测试 (ASM模式)${NC}"
  echo -e "${BLUE}========================================${NC}"
  echo "测试目录: $TEST_ROOT"
  echo "输出目录: $OUTPUT_DIR"
  echo "基线目录: $BASELINE_DIR"
  echo "测试用例总数: $total"
  echo ""
  
  SUCCESS=0
  FAILED=0
  case_num=0
  
  # 统计变量
  total_time_sum=0
  time_cases_count=0
  total_improvement_sum=0
  improvement_count=0
  total_degradation_sum=0
  degradation_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 "================================================"
    if [[ $improvement_count -gt 0 ]]; then
      avg_improvement=$(echo "scale=2; $total_improvement_sum / $improvement_count" | bc 2>/dev/null || echo "0")
      echo "平均性能提升: ${avg_improvement}% (基于 $improvement_count 个测试)"
    fi
    if [[ $degradation_count -gt 0 ]]; then
      avg_degradation=$(echo "scale=2; $total_degradation_sum / $degradation_count" | bc 2>/dev/null || echo "0")
      echo "平均性能退化: ${avg_degradation}% (基于 $degradation_count 个测试)"
    fi
    echo ""
    echo "================================================"
  } >> "$OPT_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}Lab5 批量测试完成${NC}"
  echo -e "${BLUE}执行用例: $case_num${NC}"
  echo -e "${GREEN}成功: $SUCCESS${NC}"
  echo -e "${RED}失败: $FAILED${NC}"
  echo -e "${BLUE}成功率: ${RATE}%${NC}"
  
  # 显示平均运行时间
  if [[ $time_cases_count -gt 0 ]]; then
    avg_time=$((total_time_sum / time_cases_count))
    echo -e "${BLUE}平均运行时间: ${avg_time}ms (基于 ${time_cases_count} 个成功用例)${NC}"
  fi
  
  # 显示平均优化提升
  if [[ $improvement_count -gt 0 ]]; then
    avg_improvement=$(echo "scale=2; $total_improvement_sum / $improvement_count" | bc 2>/dev/null || echo "0")
    echo -e "${GREEN}平均性能提升: ${avg_improvement}% (基于 ${improvement_count} 个测试)${NC}"
  fi
  if [[ $degradation_count -gt 0 ]]; then
    avg_degradation=$(echo "scale=2; $total_degradation_sum / $degradation_count" | bc 2>/dev/null || echo "0")
    echo -e "${RED}平均性能退化: ${avg_degradation}% (基于 ${degradation_count} 个测试)${NC}"
  fi
  
  echo -e "${BLUE}时间汇总: $TIME_SUMMARY_FILE${NC}"
  echo -e "${BLUE}优化报告: $OPT_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