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

ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
cd "$ROOT_DIR"

BUNDLE_ROOT=""
RIRS=""
OUT_ROOT=""
ROUTINATOR_ROOT="/home/yuyr/dev/rust_playground/routinator"
RPKI_CLIENT_ROOT="/home/yuyr/dev/rpki-client-9.7"
RPKI_CLIENT_BUILD_DIR="/home/yuyr/dev/rpki-client-9.7/build-m5"
KEEP_DB=0

usage() {
  cat <<'EOF'
Usage:
  ./scripts/replay_verify/run_peer_bundle_matrix.sh \
    --bundle-root <dir> \
    [--rir <afrinic,apnic,...>] \
    [--out-root <dir>] \
    [--routinator-root <dir>] \
    [--rpki-client-root <dir>] \
    [--rpki-client-build-dir <dir>] \
    [--keep-db]
EOF
}

while [[ $# -gt 0 ]]; do
  case "$1" in
    --bundle-root) BUNDLE_ROOT="${2:?}"; shift 2 ;;
    --rir) RIRS="${2:?}"; shift 2 ;;
    --out-root) OUT_ROOT="${2:?}"; shift 2 ;;
    --routinator-root) ROUTINATOR_ROOT="${2:?}"; shift 2 ;;
    --rpki-client-root) RPKI_CLIENT_ROOT="${2:?}"; shift 2 ;;
    --rpki-client-build-dir) RPKI_CLIENT_BUILD_DIR="${2:?}"; shift 2 ;;
    --keep-db) KEEP_DB=1; shift ;;
    --help|-h) usage; exit 0 ;;
    *) echo "unknown argument: $1" >&2; usage >&2; exit 2 ;;
  esac
done

if [[ -z "$BUNDLE_ROOT" ]]; then
  usage >&2
  exit 2
fi

BUNDLE_ROOT="$(python3 - "$BUNDLE_ROOT" <<'PY'
from pathlib import Path
import sys
print(Path(sys.argv[1]).resolve())
PY
)"

RUN_TAG="$(date -u +%Y%m%dT%H%M%SZ)"
if [[ -z "$OUT_ROOT" ]]; then
  OUT_ROOT="target/replay/peer_bundle_matrix_${RUN_TAG}"
fi
mkdir -p "$OUT_ROOT"
OUT_ROOT="$(python3 - "$OUT_ROOT" <<'PY'
from pathlib import Path
import sys
print(Path(sys.argv[1]).resolve())
PY
)"

discover_rirs() {
  python3 - "$BUNDLE_ROOT" <<'PY'
from pathlib import Path
import sys
root = Path(sys.argv[1])
if (root / "base-locks.json").exists():
    print(root.name)
    raise SystemExit
rirs = []
for entry in sorted(root.iterdir()):
    if not entry.is_dir():
        continue
    if (entry / "base-locks.json").exists():
        rirs.append(entry.name)
        continue
    nested = sorted(
        child.name for child in entry.iterdir()
        if child.is_dir() and (child / "base-locks.json").exists()
    )
    if len(nested) == 1:
        rirs.append(nested[0])
print(",".join(rirs))
PY
}

if [[ -z "$RIRS" ]]; then
  RIRS="$(discover_rirs)"
fi

ROUTI_OUT="$OUT_ROOT/routinator"
CLIENT_OUT="$OUT_ROOT/rpki-client"
NORMALIZED_BUNDLE_ROOT="$OUT_ROOT/.normalized-bundle-root"
mkdir -p "$ROUTI_OUT" "$CLIENT_OUT"
rm -rf "$NORMALIZED_BUNDLE_ROOT"
mkdir -p "$NORMALIZED_BUNDLE_ROOT"

IFS=',' read -r -a RIR_LIST <<< "$RIRS"
for raw_rir in "${RIR_LIST[@]}"; do
  rir="$(printf '%s' "$raw_rir" | tr '[:upper:]' '[:lower:]' | xargs)"
  [[ -n "$rir" ]] || continue
  source_bundle_dir=""
  if [[ -d "$BUNDLE_ROOT/$rir" && -f "$BUNDLE_ROOT/$rir/base-locks.json" ]]; then
    source_bundle_dir="$BUNDLE_ROOT/$rir"
  else
    match="$(find "$BUNDLE_ROOT" -maxdepth 2 -type d -path "*/${rir}" -exec test -f '{}/base-locks.json' ';' -print | head -n 1)"
    if [[ -z "$match" ]]; then
      echo "unable to resolve bundle directory for RIR: $rir" >&2
      exit 1
    fi
    source_bundle_dir="$match"
  fi
  ln -sfn "$source_bundle_dir" "$NORMALIZED_BUNDLE_ROOT/$rir"
  ROUTI_CMD=(
    "$ROUTINATOR_ROOT/bench/multi_rir_demo_ours/run_single_rir_ours_bundle.sh"
    "$source_bundle_dir"
    "$ROUTI_OUT/$rir"
  )
  [[ "$KEEP_DB" -eq 1 ]] && ROUTI_CMD=( "$ROUTINATOR_ROOT/bench/multi_rir_demo_ours/run_single_rir_ours_bundle.sh" --keep-db "$source_bundle_dir" "$ROUTI_OUT/$rir" )
  "${ROUTI_CMD[@]}"
done

CLIENT_ARGS=(
  python3 "$RPKI_CLIENT_ROOT/tools/run_bundle_matrix.py"
  --bundle-dir "$NORMALIZED_BUNDLE_ROOT"
  --build-dir "$RPKI_CLIENT_BUILD_DIR"
  --work-dir "$CLIENT_OUT"
)
[[ "$KEEP_DB" -eq 1 ]] && CLIENT_ARGS+=(--keep-db)
for raw_rir in "${RIR_LIST[@]}"; do
  rir="$(printf '%s' "$raw_rir" | tr '[:upper:]' '[:lower:]' | xargs)"
  [[ -n "$rir" ]] || continue
  CLIENT_ARGS+=(--rir "$rir")
done
"${CLIENT_ARGS[@]}"

SUMMARY_JSON="$OUT_ROOT/summary.json"
SUMMARY_MD="$OUT_ROOT/summary.md"

python3 - "$ROUTI_OUT" "$CLIENT_OUT/matrix-summary.json" "$SUMMARY_JSON" <<'PY'
import json
from pathlib import Path
import sys

routi_root = Path(sys.argv[1])
client_summary = json.loads(Path(sys.argv[2]).read_text())
summary_path = Path(sys.argv[3])
summary = {"routinator": {}, "rpki_client": client_summary}
for verification in sorted(routi_root.glob("*/verification.json")):
    rir = verification.parent.name
    summary["routinator"][rir] = json.loads(verification.read_text())
summary_path.write_text(json.dumps(summary, indent=2))
PY

python3 - "$SUMMARY_JSON" "$SUMMARY_MD" <<'PY'
import json
from pathlib import Path
import sys
summary = json.loads(Path(sys.argv[1]).read_text())
out = Path(sys.argv[2])
lines = [
    "# Peer Bundle Matrix Summary",
    "",
    "## Routinator",
    "",
    "| rir | base_vrp | base_vap | sequence_vrp | sequence_vap |",
    "|---|---|---|---|---|",
]
for rir, data in sorted(summary["routinator"].items()):
    if "steps" in data:
        lines.append(
            f"| {rir} | {str(data.get('baseMatch')).lower()} | {str(data.get('baseVapsMatch')).lower()} | "
            f"{str(data.get('summary', {}).get('allStepsMatch')).lower()} | "
            f"{str(data.get('summary', {}).get('allStepsVapsMatch')).lower()} |"
        )
    else:
        lines.append(
            f"| {rir} | {str(data.get('baseMatch')).lower()} | {str(data.get('baseVapsMatch')).lower()} | "
            f"{str(data.get('deltaMatch')).lower()} | {str(data.get('deltaVapsMatch')).lower()} |"
        )
lines += [
    "",
    "## rpki-client",
    "",
    "| rir | base_vrp | base_vap | sequence_vrp | sequence_vap |",
    "|---|---|---|---|---|",
]
for rir, phases in sorted(summary["rpki_client"].items()):
    base = phases.get("base", {})
    step_items = [
        value for key, value in phases.items()
        if key not in ("base", "delta") and isinstance(value, dict)
    ]
    if "delta" in phases:
        step_items.append(phases["delta"])
    all_step_match = all(item.get("match") for item in step_items) if step_items else None
    all_step_vap_match = all(item.get("vaps_match") for item in step_items) if step_items else None
    lines.append(
        f"| {rir} | {str(base.get('match')).lower()} | {str(base.get('vaps_match')).lower()} | "
        f"{str(all_step_match).lower()} | {str(all_step_vap_match).lower()} |"
    )
out.write_text("\n".join(lines) + "\n")
PY

echo "$OUT_ROOT"