#!/usr/bin/env bash
set -euo pipefail

ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
FIO_CASE_DIR="${ROOT_DIR}/zvfs_fio_test"
RESULTS_ROOT="${ROOT_DIR}/results/fio"
TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
RUN_DIR="${RESULTS_ROOT}/${TIMESTAMP}"

FIO_BIN="${FIO_BIN:-$(command -v fio 2>/dev/null || true)}"
PYTHON_BIN="${PYTHON_BIN:-$(command -v python3 2>/dev/null || true)}"
LD_PRELOAD_PATH="${LD_PRELOAD_PATH:-${ROOT_DIR}/src/libzvfs.so}"
ZVFS_LD_PRELOAD_VALUE="${ZVFS_LD_PRELOAD_VALUE:-${LD_PRELOAD_PATH}}"

SYS_PARENT_DIR="/tmp/zvfs_sys_fio"
ZVFS_PARENT_DIR="/zvfs/zvfs_sys_fio"
SYS_TEST_FILE="${SYS_PARENT_DIR}/testfile"
ZVFS_TEST_FILE="${ZVFS_PARENT_DIR}/testfile"

WORKLOADS=(
  "prepare_fill"
  "randread_4k"
  "randwrite_4k"
  "randrw_4k"
)

MODES=(
  "sys"
  "zvfs"
)

require_bin() {
  local bin_path="$1"
  local name="$2"

  if [[ -z "${bin_path}" || ! -x "${bin_path}" ]]; then
    echo "未找到 ${name}，请先安装或通过环境变量指定路径。" >&2
    exit 1
  fi
}

ensure_paths() {
  mkdir -p "${SYS_PARENT_DIR}"

  if [[ ! -d "${ZVFS_PARENT_DIR}" ]]; then
    echo "未找到 ZVFS 测试目录：${ZVFS_PARENT_DIR}" >&2
    echo "请先准备好对应目录，再执行脚本。" >&2
    exit 1
  fi

  if [[ -z "${ZVFS_LD_PRELOAD_VALUE}" ]]; then
    echo "ZVFS_LD_PRELOAD_VALUE 不能为空。" >&2
    exit 1
  fi

  local zvfs_preload_file="${ZVFS_LD_PRELOAD_VALUE%%:*}"
  if [[ ! -f "${zvfs_preload_file}" ]]; then
    echo "未找到 libzvfs.so：${zvfs_preload_file}" >&2
    exit 1
  fi
}

cleanup_test_file() {
  local mode="$1"

  if [[ "${mode}" == "zvfs" ]]; then
    env LD_PRELOAD="${ZVFS_LD_PRELOAD_VALUE}" rm -f "${ZVFS_TEST_FILE}"
  else
    rm -f "${SYS_TEST_FILE}"
  fi
}

cleanup_all_test_files() {
  cleanup_test_file "sys"
  cleanup_test_file "zvfs"
}

run_case() {
  local mode="$1"
  local workload="$2"
  local fio_file="${FIO_CASE_DIR}/${mode}/${workload}.fio"
  local output_file="${RUN_DIR}/${mode}_${workload}.json"

  if [[ ! -f "${fio_file}" ]]; then
    echo "未找到 fio 配置：${fio_file}" >&2
    exit 1
  fi

  echo "[run] mode=${mode} workload=${workload}"

  if [[ "${workload}" == "prepare_fill" ]]; then
    cleanup_test_file "${mode}"
  fi

  if [[ "${mode}" == "zvfs" ]]; then
    env LD_PRELOAD="${ZVFS_LD_PRELOAD_VALUE}" \
      "${FIO_BIN}" --output-format=json --output="${output_file}" "${fio_file}"
  else
    "${FIO_BIN}" --output-format=json --output="${output_file}" "${fio_file}"
  fi
}

write_summary() {
  local summary_file="${RUN_DIR}/summary.md"
  local parser_script
  parser_script="$(cat <<'PY'
import json
import os
import sys

run_dir = sys.argv[1]
modes = ["sys", "zvfs"]
workloads = ["prepare_fill", "randread_4k", "randwrite_4k", "randrw_4k"]

def load_job(path):
    with open(path, "r", encoding="utf-8") as f:
        content = f.read()
    start = content.find("{")
    if start == -1:
        raise RuntimeError(f"fio output is not JSON: {path}")
    data = json.loads(content[start:])
    jobs = data.get("jobs", [])
    if not jobs:
        raise RuntimeError(f"fio output has no jobs: {path}")
    return data, jobs[0]

def mib_per_sec(io_part):
    value = io_part.get("bw_bytes", 0)
    if not value:
        return "-"
    return f"{value / 1024 / 1024:.2f}"

def iops(io_part):
    value = io_part.get("iops", 0)
    if not value:
        return "-"
    return f"{value:.2f}"

def lat_us(io_part):
    clat_ns = io_part.get("clat_ns", {})
    value = clat_ns.get("mean", 0)
    if not value:
        return "-"
    return f"{value / 1000:.2f}"

def disk_util(data):
    parts = []
    for item in data.get("disk_util", []):
        name = item.get("name")
        util = item.get("util")
        if not name or util is None:
            continue
        parts.append(f"{name}={util:.2f}%")
    return ", ".join(parts) if parts else "-"

lines = [
    "# FIO Summary",
    "",
    f"run_dir: `{run_dir}`",
    "",
    "| mode | workload | bw(MiB/s) | iops | avg_lat(us) | read_iops | write_iops | read_avg_lat(us) | write_avg_lat(us) | disk_util |",
    "| --- | --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | --- |",
]

for mode in modes:
    for workload in workloads:
        path = os.path.join(run_dir, f"{mode}_{workload}.json")
        data, job = load_job(path)
        read_part = job.get("read", {})
        write_part = job.get("write", {})
        util_col = disk_util(data)

        if workload == "prepare_fill":
            row = [
                mode,
                workload,
                mib_per_sec(write_part),
                "-",
                "-",
                "-",
                "-",
                "-",
                "-",
                util_col,
            ]
        elif workload == "randread_4k":
            row = [
                mode,
                workload,
                mib_per_sec(read_part),
                iops(read_part),
                lat_us(read_part),
                "-",
                "-",
                "-",
                "-",
                util_col,
            ]
        elif workload == "randwrite_4k":
            row = [
                mode,
                workload,
                mib_per_sec(write_part),
                iops(write_part),
                lat_us(write_part),
                "-",
                "-",
                "-",
                "-",
                util_col,
            ]
        else:
            row = [
                mode,
                workload,
                "-",
                "-",
                "-",
                iops(read_part),
                iops(write_part),
                lat_us(read_part),
                lat_us(write_part),
                util_col,
            ]

        lines.append("| " + " | ".join(row) + " |")

print("\n".join(lines))
PY
)"

  "${PYTHON_BIN}" -c "${parser_script}" "${RUN_DIR}" > "${summary_file}"
  echo
  echo "[summary] ${summary_file}"
  sed -n '1,200p' "${summary_file}"
}

main() {
  require_bin "${FIO_BIN}" "fio"
  require_bin "${PYTHON_BIN}" "python3"
  ensure_paths

  mkdir -p "${RUN_DIR}"

  echo "结果目录：${RUN_DIR}"
  echo "fio 配置目录：${FIO_CASE_DIR}"
  echo "LD_PRELOAD：${ZVFS_LD_PRELOAD_VALUE}"
  echo

  cleanup_all_test_files

  for mode in "${MODES[@]}"; do
    for workload in "${WORKLOADS[@]}"; do
      run_case "${mode}" "${workload}"
    done
  done

  write_summary
  cleanup_all_test_files
}

main "$@"