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

usage() {
  cat <<'EOF'
Usage:
  ./scripts/cir/run_cir_replay_ours.sh \
    --cir <path> \
    --repo-bytes-db <path> \
    --out-dir <path> \
    --reference-ccr <path> \
    [--keep-db] \
    [--write-actual-ccr] \
    [--write-report-json] \
    [--report-json-compact] \
    [--phase2-object-workers <n>] \
    [--phase2-worker-queue-capacity <n>] \
    [--rpki-bin <path>] \
    [--real-rsync-bin <path>]
EOF
}

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

CIR=""
REPO_BYTES_DB=""
OUT_DIR=""
REFERENCE_CCR=""
KEEP_DB=0
WRITE_ACTUAL_CCR=0
WRITE_REPORT_JSON=0
REPORT_JSON_COMPACT=0
PHASE2_OBJECT_WORKERS="${CIR_REPLAY_PHASE2_OBJECT_WORKERS:-4}"
PHASE2_WORKER_QUEUE_CAPACITY="${CIR_REPLAY_PHASE2_WORKER_QUEUE_CAPACITY:-64}"
RPKI_BIN="${RPKI_BIN:-$ROOT_DIR/target/release/rpki}"
CIR_MATERIALIZE_BIN="${CIR_MATERIALIZE_BIN:-$ROOT_DIR/target/release/cir_materialize}"
CIR_EXTRACT_INPUTS_BIN="${CIR_EXTRACT_INPUTS_BIN:-$ROOT_DIR/target/release/cir_extract_inputs}"
CCR_TO_COMPARE_VIEWS_BIN="${CCR_TO_COMPARE_VIEWS_BIN:-$ROOT_DIR/target/release/ccr_to_compare_views}"
REAL_RSYNC_BIN="${REAL_RSYNC_BIN:-/usr/bin/rsync}"
WRAPPER="$ROOT_DIR/scripts/cir/cir-rsync-wrapper"

while [[ $# -gt 0 ]]; do
  case "$1" in
    --cir) CIR="$2"; shift 2 ;;
    --repo-bytes-db) REPO_BYTES_DB="$2"; shift 2 ;;
    --out-dir) OUT_DIR="$2"; shift 2 ;;
    --reference-ccr) REFERENCE_CCR="$2"; shift 2 ;;
    --keep-db) KEEP_DB=1; shift ;;
    --write-actual-ccr) WRITE_ACTUAL_CCR=1; shift ;;
    --write-report-json) WRITE_REPORT_JSON=1; shift ;;
    --report-json-compact) WRITE_REPORT_JSON=1; REPORT_JSON_COMPACT=1; shift ;;
    --phase2-object-workers) PHASE2_OBJECT_WORKERS="$2"; shift 2 ;;
    --phase2-worker-queue-capacity) PHASE2_WORKER_QUEUE_CAPACITY="$2"; shift 2 ;;
    --rpki-bin) RPKI_BIN="$2"; shift 2 ;;
    --real-rsync-bin) REAL_RSYNC_BIN="$2"; shift 2 ;;
    -h|--help) usage; exit 0 ;;
    *) echo "unknown argument: $1" >&2; usage; exit 2 ;;
  esac
done

[[ -n "$CIR" && -n "$REPO_BYTES_DB" && -n "$OUT_DIR" && -n "$REFERENCE_CCR" ]] || {
  usage >&2
  exit 2
}

mkdir -p "$OUT_DIR"

needs_build=0
if [[ ! -x "$RPKI_BIN" || ! -x "$CIR_MATERIALIZE_BIN" || ! -x "$CIR_EXTRACT_INPUTS_BIN" || ! -x "$CCR_TO_COMPARE_VIEWS_BIN" ]]; then
  needs_build=1
elif [[ "$RPKI_BIN" == "$ROOT_DIR/target/release/rpki" ]] && find "$ROOT_DIR/src" "$ROOT_DIR/Cargo.toml" -newer "$RPKI_BIN" -print -quit | grep -q .; then
  needs_build=1
fi
if [[ "$needs_build" -eq 1 ]]; then
  (
    cd "$ROOT_DIR"
    cargo build --release --bin rpki --bin cir_materialize --bin cir_extract_inputs --bin ccr_to_compare_views
  )
fi

TMP_ROOT="$OUT_DIR/.tmp"
TALS_DIR="$TMP_ROOT/tals"
META_JSON="$TMP_ROOT/meta.json"
MIRROR_ROOT="$TMP_ROOT/mirror"
DB_DIR="$TMP_ROOT/work-db"
REPLAY_RAW_STORE_DB="$TMP_ROOT/replay-raw-store.db"
REPLAY_REPO_BYTES_DB="$TMP_ROOT/replay-repo-bytes.db"
ACTUAL_CCR="$OUT_DIR/actual.ccr"
ACTUAL_REPORT="$OUT_DIR/report.json"
ACTUAL_VRPS="$OUT_DIR/actual-vrps.csv"
ACTUAL_VAPS="$OUT_DIR/actual-vaps.csv"
REF_VRPS="$OUT_DIR/reference-vrps.csv"
REF_VAPS="$OUT_DIR/reference-vaps.csv"
COMPARE_JSON="$OUT_DIR/compare-summary.json"
RUN_LOG="$OUT_DIR/run.log"

rm -rf "$TMP_ROOT"
mkdir -p "$TMP_ROOT"

"$CIR_EXTRACT_INPUTS_BIN" --cir "$CIR" --tals-dir "$TALS_DIR" --meta-json "$META_JSON"
materialize_cmd=("$CIR_MATERIALIZE_BIN" --cir "$CIR" --repo-bytes-db "$REPO_BYTES_DB" --mirror-root "$MIRROR_ROOT")
if [[ "$KEEP_DB" -eq 1 ]]; then
  materialize_cmd+=(--keep-db)
fi
"${materialize_cmd[@]}"

VALIDATION_TIME="$(python3 - <<'PY' "$META_JSON"
import json,sys
print(json.load(open(sys.argv[1]))["validationTime"])
PY
)"
mapfile -t TAL_PATHS < <(python3 - <<'PY' "$META_JSON"
import json, sys
for item in json.load(open(sys.argv[1], encoding="utf-8"))["talFiles"]:
    print(item["path"])
PY
)
TAL_ARGS=()
for tal_path in "${TAL_PATHS[@]}"; do
  TAL_ARGS+=(--tal-path "$tal_path")
done

export CIR_MIRROR_ROOT="$(python3 - <<'PY' "$MIRROR_ROOT"
from pathlib import Path
import sys
print(Path(sys.argv[1]).resolve())
PY
)"
export REAL_RSYNC_BIN="$REAL_RSYNC_BIN"
export CIR_LOCAL_LINK_MODE=1

REPORT_JSON_ARGS=(--skip-report-build)
VCIR_ARGS=(--skip-vcir-persist)
if [[ "$WRITE_REPORT_JSON" -eq 1 ]]; then
  REPORT_JSON_ARGS=(--report-json "$ACTUAL_REPORT")
  if [[ "$REPORT_JSON_COMPACT" -eq 1 ]]; then
    REPORT_JSON_ARGS+=(--report-json-compact)
  fi
fi
CCR_ARGS=()
if [[ "$WRITE_ACTUAL_CCR" -eq 1 ]]; then
  CCR_ARGS=(--ccr-out "$ACTUAL_CCR")
fi

"$RPKI_BIN" \
  --db "$DB_DIR" \
  --raw-store-db "$REPLAY_RAW_STORE_DB" \
  --repo-bytes-db "$REPLAY_REPO_BYTES_DB" \
  "${TAL_ARGS[@]}" \
  --parallel-phase2-object-workers "$PHASE2_OBJECT_WORKERS" \
  --parallel-phase2-worker-queue-capacity "$PHASE2_WORKER_QUEUE_CAPACITY" \
  --disable-rrdp \
  --rsync-command "$WRAPPER" \
  --validation-time "$VALIDATION_TIME" \
  "${CCR_ARGS[@]}" \
  --vrps-csv-out "$ACTUAL_VRPS" \
  --vaps-csv-out "$ACTUAL_VAPS" \
  --compare-view-trust-anchor unknown \
  "${VCIR_ARGS[@]}" \
  "${REPORT_JSON_ARGS[@]}" \
  >"$RUN_LOG" 2>&1

sort_compare_csv() {
  local path="$1"
  local tmp="${path}.sorted.tmp"
  {
    head -n 1 "$path"
    tail -n +2 "$path" | LC_ALL=C sort -u
  } >"$tmp"
  mv "$tmp" "$path"
}

sort_compare_csv "$ACTUAL_VRPS"
sort_compare_csv "$ACTUAL_VAPS"
"$CCR_TO_COMPARE_VIEWS_BIN" --ccr "$REFERENCE_CCR" --vrps-out "$REF_VRPS" --vaps-out "$REF_VAPS" --trust-anchor unknown

python3 - <<'PY' "$ACTUAL_VRPS" "$REF_VRPS" "$ACTUAL_VAPS" "$REF_VAPS" "$COMPARE_JSON" "$META_JSON" "$WRITE_REPORT_JSON" "$WRITE_ACTUAL_CCR" "$PHASE2_OBJECT_WORKERS" "$PHASE2_WORKER_QUEUE_CAPACITY"
import csv, json, sys

def next_row(reader):
    try:
        return tuple(next(reader))
    except StopIteration:
        return None

def compare_sorted_csv(actual_path, ref_path):
    actual_count = 0
    ref_count = 0
    only_actual_count = 0
    only_ref_count = 0
    only_actual_sample = []
    only_ref_sample = []
    with open(actual_path, newline="") as actual_file, open(ref_path, newline="") as ref_file:
        actual_reader = csv.reader(actual_file)
        ref_reader = csv.reader(ref_file)
        next(actual_reader, None)
        next(ref_reader, None)
        actual = next_row(actual_reader)
        ref = next_row(ref_reader)
        while actual is not None or ref is not None:
            if ref is None or (actual is not None and actual < ref):
                actual_count += 1
                only_actual_count += 1
                if len(only_actual_sample) < 20:
                    only_actual_sample.append(list(actual))
                actual = next_row(actual_reader)
            elif actual is None or ref < actual:
                ref_count += 1
                only_ref_count += 1
                if len(only_ref_sample) < 20:
                    only_ref_sample.append(list(ref))
                ref = next_row(ref_reader)
            else:
                actual_count += 1
                ref_count += 1
                actual = next_row(actual_reader)
                ref = next_row(ref_reader)
    return {
        "actual": actual_count,
        "reference": ref_count,
        "only_in_actual": only_actual_sample,
        "only_in_reference": only_ref_sample,
        "only_in_actual_count": only_actual_count,
        "only_in_reference_count": only_ref_count,
        "match": only_actual_count == 0 and only_ref_count == 0,
    }

vrps = compare_sorted_csv(sys.argv[1], sys.argv[2])
vaps = compare_sorted_csv(sys.argv[3], sys.argv[4])
meta = json.load(open(sys.argv[6], encoding="utf-8"))
summary = {
  "compareMode": "trust-anchor-agnostic",
  "talCount": len(meta["talFiles"]),
  "talPaths": [item["path"] for item in meta["talFiles"]],
  "actualCcrWritten": sys.argv[8] == "1",
  "reportJsonWritten": sys.argv[7] == "1",
  "replayParallelism": {
    "phase2ObjectWorkers": int(sys.argv[9]),
    "phase2WorkerQueueCapacity": int(sys.argv[10]),
  },
  "vrps": vrps,
  "vaps": vaps,
}
with open(sys.argv[5], "w") as f:
    json.dump(summary, f, indent=2)
PY

if [[ "$KEEP_DB" -ne 1 ]]; then
  rm -rf "$TMP_ROOT"
fi

echo "done: $OUT_DIR"