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

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
AI_RESEARCH_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"
TIMESTAMP="${ENTRYPOINT_CONTRACT_TIMESTAMP:-$(date -u +%Y%m%dT%H%M%SZ)}"
PARITY_DIR="${AI_RESEARCH_ROOT}/benchmarks/parity_logs/entrypoint_contract_${TIMESTAMP}"
AUDIT_DIR="${AI_RESEARCH_ROOT}/audit/entrypoint_contract_${TIMESTAMP}"

mkdir -p "$PARITY_DIR" "$AUDIT_DIR"

log() {
  printf '%s\n' "$*" | tee -a "$PARITY_DIR/contract_check.log"
}

fail() {
  echo "ERROR: $*" | tee -a "$PARITY_DIR/contract_check.log"
}

resolve_entrypoint() {
  local role="$1"
  local override_name="$2"
  local pattern="$3"
  local -a candidates
  local override="${!override_name:-}"

  if [[ -n "$override" ]]; then
    if [[ -x "$override" ]]; then
      echo "$override"
      return 0
    fi
    fail "${role} override is not executable: $override"
    return 1
  fi

  candidates=()
  if [[ "$role" == "canonical" ]]; then
    candidates+=("${AI_RESEARCH_ROOT}/benchmarks/verify_benchmark_artifacts.sh")
  fi
  if [[ "$role" == "wrapper" ]]; then
    candidates+=("${AI_RESEARCH_ROOT}/benchmarks/verify_benchmark_artifacts_wrapper.sh")
  fi
  candidates+=(
    "${AI_RESEARCH_ROOT}/benchmarks/${role}_entrypoint.sh"
    "${AI_RESEARCH_ROOT}/benchmarks/entrypoint_${role}.sh"
    "${AI_RESEARCH_ROOT}/benchmarks/run_${role}_entrypoint.sh"
    "${AI_RESEARCH_ROOT}/benchmarks/${role}.sh"
    "${AI_RESEARCH_ROOT}/benchmarks/run_${role}.sh"
  )

  local c
  for c in "${candidates[@]}"; do
    if [[ -x "$c" ]]; then
      echo "$c"
      return 0
    fi
  done

  # Fallback scan for likely files when naming differs
  while IFS= read -r c; do
    if [[ -x "$c" ]]; then
      echo "$c"
      return 0
    fi
  done < <(find "${AI_RESEARCH_ROOT}" -maxdepth 6 -type f -name "*.sh" \
      \( -iname "*${pattern}*" -o -iname "*${role}*entrypoint*" -o -iname "*verify*artifact*${role}*" \) \
      | sort)

  return 1
}

run_entrypoint() {
  local tag="$1"
  local entrypoint="$2"
  shift 2
  local -a args=("$@")
  local out="${PARITY_DIR}/${tag}.log"
  local code_file="${PARITY_DIR}/${tag}.code"

  set +e
  "${entrypoint}" "${args[@]}" >"$out" 2>&1
  local rc=$?
  set -e

  printf '%s\n' "$rc" >"$code_file"
  echo "$rc"
}

CANONICAL_OVERRIDE="${CANONICAL_ENTRYPOINT:-}"
WRAPPER_OVERRIDE="${WRAPPER_ENTRYPOINT:-}"
CANONICAL="$(resolve_entrypoint canonical CANONICAL_ENTRYPOINT 'canonical_entrypoint')" || CANONICAL=""
WRAPPER="$(resolve_entrypoint wrapper WRAPPER_ENTRYPOINT 'wrapper_entrypoint')" || WRAPPER=""

if [[ -z "$CANONICAL" || -z "$WRAPPER" ]]; then
  fail "Failed to resolve canonical/wrapper entrypoints."
  fail "canonical=${CANONICAL:-missing} wrapper=${WRAPPER:-missing}"
  fail "Set CANONICAL_ENTRYPOINT and WRAPPER_ENTRYPOINT explicitly if auto-discovery is insufficient."
  exit 1
fi

if [[ ! -x "$CANONICAL" || ! -x "$WRAPPER" ]]; then
  fail "Resolved entrypoints must be executable."
  fail "canonical=${CANONICAL} wrapper=${WRAPPER}"
  exit 1
fi

log "Canonical entrypoint: $CANONICAL"
log "Wrapper entrypoint:   $WRAPPER"

# newline-separated list; default is a single smoke scenario to avoid long benchmark runs
SCENARIO_SPEC="${ENTRYPOINT_SCENARIOS:-$'--help'}"
mapfile -t SCENARIOS < <(printf '%s' "$SCENARIO_SPEC" | tr -d '\r' | sed '/^$/d')
if (( ${#SCENARIOS[@]} == 0 )); then
  SCENARIOS=("")
fi

FAILED=0
SCENARIO_INDEX=1
for scenario in "${SCENARIOS[@]}"; do
  read -r -a scenario_args <<< "$scenario"

  CAN_LABEL="scenario${SCENARIO_INDEX}_canonical"
  WRAP_LABEL="scenario${SCENARIO_INDEX}_wrapper"

  can_rc="$(run_entrypoint "$CAN_LABEL" "$CANONICAL" "${scenario_args[@]}")"
  wrap_rc="$(run_entrypoint "$WRAP_LABEL" "$WRAPPER" "${scenario_args[@]}")"

  if [[ "$can_rc" != "$wrap_rc" ]]; then
    fail "Exit-code mismatch for scenario ${SCENARIO_INDEX}: '$scenario' canonical=$can_rc wrapper=$wrap_rc"
    FAILED=1
  fi

  if ! diff -u "$PARITY_DIR/${CAN_LABEL}.log" "$PARITY_DIR/${WRAP_LABEL}.log" > "$PARITY_DIR/scenario${SCENARIO_INDEX}.diff"; then
    fail "Output mismatch for scenario ${SCENARIO_INDEX}: '$scenario'"
    FAILED=1
  fi
  SCENARIO_INDEX=$((SCENARIO_INDEX + 1))
done

{
  printf '%s\n' "entrypoint_contract_check"
  printf 'timestamp=%s\n' "$TIMESTAMP"
  printf 'canonical=%s\n' "$CANONICAL"
  printf 'wrapper=%s\n' "$WRAPPER"
  printf 'scenarios=%s\n' "${#SCENARIOS[@]}"
  if (( FAILED == 0 )); then
    printf 'status=PASS\n'
  else
    printf 'status=FAIL\n'
  fi
} > "$PARITY_DIR/summary.txt"
cp "$PARITY_DIR/summary.txt" "$AUDIT_DIR/summary.txt"
cp "$PARITY_DIR"/*.log "$PARITY_DIR"/*.code "$PARITY_DIR"/*.diff "$AUDIT_DIR"/ 2>/dev/null || true

if (( FAILED != 0 )); then
  log "Parity check failed; see audit=$AUDIT_DIR and parity=$PARITY_DIR"
  exit 1
fi

log "Parity check passed; evidence written to audit=$AUDIT_DIR and parity=$PARITY_DIR"
exit 0