#!/usr/bin/env bash
# Model Candidate Evaluation Script
# Phase 1: Smoke test (load model, embed one query) — ~30s per model
# Phase 2: Val benchmark (full real-v1-val dataset) — ~20-80min per model
#
# Usage:
#   ./scripts/eval-candidates.sh              # Run all phases
#   ./scripts/eval-candidates.sh --smoke-only  # Phase 1 only (fast)
#   ./scripts/eval-candidates.sh --bench-only  # Phase 2 only (skip smoke, assume all pass)
#   ./scripts/eval-candidates.sh --model gte-modernbert-base  # Test single model

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
PROJECT_DIR="$(cd "$SCRIPT_DIR/.." && pwd)"
cd "$PROJECT_DIR"

RESULTS_DIR="$PROJECT_DIR/.bluera/bluera-knowledge/bench-data/model-eval"
mkdir -p "$RESULTS_DIR"

SMOKE_ONLY=false
BENCH_ONLY=false
SINGLE_MODEL=""

while [[ $# -gt 0 ]]; do
  case $1 in
    --smoke-only) SMOKE_ONLY=true; shift ;;
    --bench-only) BENCH_ONLY=true; shift ;;
    --model) SINGLE_MODEL="$2"; shift 2 ;;
    *) echo "Unknown option: $1"; exit 1 ;;
  esac
done

# Priority-ordered candidate list
CANDIDATES=(
  "gte-modernbert-base"
  "snowflake-arctic-embed-s"
  "snowflake-arctic-embed-m-v1.5"
  "jina-embeddings-v2-base-code"
  "modernbert-embed-base"
  "snowflake-arctic-embed-xs"
  "snowflake-arctic-embed-m-v2.0"
)

if [[ -n "$SINGLE_MODEL" ]]; then
  CANDIDATES=("$SINGLE_MODEL")
fi

SMOKE_PASSED=()
SMOKE_FAILED=()

# ============================================================
# Phase 1: Smoke Tests
# ============================================================
smoke_test() {
  local model_key="$1"
  local result_file="$RESULTS_DIR/smoke-${model_key}.json"

  echo -n "  [$model_key] Loading model... "

  # Use bun to run inline TypeScript that loads model via our pipeline
  local output
  if output=$(timeout 120 bun -e "
import { MODEL_REGISTRY } from './src/models/registry.ts';
import { pipeline } from '@huggingface/transformers';

const key = '${model_key}';
const config = MODEL_REGISTRY[key];
if (!config) {
  console.error('Model not in registry: ' + key);
  process.exit(1);
}

console.error('Downloading ' + config.id + '...');
const start = Date.now();

try {
  const extractor = await pipeline('feature-extraction', config.id, { dtype: 'fp32' });

  const testQuery = config.queryPrefix + 'How do I implement dependency injection?';
  const result = await extractor(testQuery, { pooling: config.pooling, normalize: config.normalize });

  const dims = result.dims;
  const elapsed = Date.now() - start;
  const sample = Array.from(result.data).slice(0, 3).map((v: number) => v.toFixed(4));

  const report = {
    model: key,
    hfId: config.id,
    status: 'pass',
    dims: dims,
    expectedDims: config.dimensions,
    dimsMatch: JSON.stringify(dims) === JSON.stringify([1, config.dimensions]),
    pooling: config.pooling,
    sampleEmbedding: sample,
    loadTimeMs: elapsed,
    timestamp: new Date().toISOString(),
  };
  console.log(JSON.stringify(report));

  await extractor.dispose();
} catch (err: any) {
  const report = {
    model: key,
    hfId: config.id,
    status: 'fail',
    error: err.message || String(err),
    timestamp: new Date().toISOString(),
  };
  console.log(JSON.stringify(report));
  process.exit(1);
}
" 2>"$RESULTS_DIR/smoke-${model_key}.log"); then
    echo "$output" > "$result_file"
    local dims load_time
    dims=$(echo "$output" | bun -e "const d=JSON.parse(await Bun.stdin.text()); process.stdout.write(String(d.dims))")
    load_time=$(echo "$output" | bun -e "const d=JSON.parse(await Bun.stdin.text()); process.stdout.write(String(d.loadTimeMs))")
    echo "PASS (dims=${dims}, ${load_time}ms)"
    return 0
  else
    echo "FAIL"
    if [[ -f "$RESULTS_DIR/smoke-${model_key}.log" ]]; then
      echo "    Error: $(tail -3 "$RESULTS_DIR/smoke-${model_key}.log")"
    fi
    return 1
  fi
}

if [[ "$BENCH_ONLY" == "false" ]]; then
  echo "========================================"
  echo "Phase 1: Smoke Tests"
  echo "========================================"
  echo ""

  for model in "${CANDIDATES[@]}"; do
    if smoke_test "$model"; then
      SMOKE_PASSED+=("$model")
    else
      SMOKE_FAILED+=("$model")
    fi
  done

  echo ""
  echo "----------------------------------------"
  echo "Smoke Test Summary"
  echo "----------------------------------------"
  echo "Passed: ${#SMOKE_PASSED[@]}/${#CANDIDATES[@]}"
  for m in "${SMOKE_PASSED[@]}"; do echo "  + $m"; done
  if [[ ${#SMOKE_FAILED[@]} -gt 0 ]]; then
    echo "Failed: ${#SMOKE_FAILED[@]}"
    for m in "${SMOKE_FAILED[@]}"; do echo "  - $m"; done
  fi
  echo ""

  if [[ "$SMOKE_ONLY" == "true" ]]; then
    echo "Done (--smoke-only). To benchmark passing models:"
    echo "  ./scripts/eval-candidates.sh --bench-only"
    exit 0
  fi
else
  # bench-only mode: assume all candidates pass smoke
  SMOKE_PASSED=("${CANDIDATES[@]}")
fi

# ============================================================
# Phase 2: Val Benchmarks
# ============================================================
if [[ ${#SMOKE_PASSED[@]} -eq 0 ]]; then
  echo "No models passed smoke test. Nothing to benchmark."
  exit 1
fi

echo "========================================"
echo "Phase 2: Val Benchmarks (real-v1-val)"
echo "========================================"
echo "Models to benchmark: ${#SMOKE_PASSED[@]}"
echo ""

BENCH_RESULTS=()

for model in "${SMOKE_PASSED[@]}"; do
  echo "========================================"
  echo "Benchmarking: $model"
  echo "========================================"

  local_artifact="$RESULTS_DIR/bench-${model}.json"

  if BK_MODEL="$model" bun run bench:search \
    --dataset real-v1-val \
    --setup --force \
    --artifacts "$local_artifact" 2>&1 | tee "$RESULTS_DIR/bench-${model}.log"; then
    BENCH_RESULTS+=("$model")
    echo ""
    echo "  -> Artifact: $local_artifact"
  else
    echo "  -> BENCHMARK FAILED for $model"
  fi
  echo ""
done

# ============================================================
# Phase 3: Comparison Summary
# ============================================================
echo "========================================"
echo "COMPARISON SUMMARY"
echo "========================================"
echo ""

# Extract key metrics from each artifact
printf "%-35s %8s %8s %8s %8s %8s\n" "Model" "Hit@1" "MRR" "nDCG@10" "R@10" "P95ms"
printf "%-35s %8s %8s %8s %8s %8s\n" "---" "---" "---" "---" "---" "---"

for model in "${BENCH_RESULTS[@]}"; do
  artifact="$RESULTS_DIR/bench-${model}.json"
  if [[ -f "$artifact" ]]; then
    bun -e "
const data = JSON.parse(await Bun.file('${artifact}').text());
const s = data.summary;
const name = '${model}'.padEnd(35);
const hit1 = (s.hitAt1 * 100).toFixed(1).padStart(7) + '%';
const mrr = s.mrr.toFixed(3).padStart(8);
const ndcg = (s.ndcgAt10 * 100).toFixed(1).padStart(7) + '%';
const r10 = (s.recallAt10 * 100).toFixed(1).padStart(7) + '%';
const p95 = s.latency.p95.toFixed(0).padStart(5) + 'ms';
console.log(name + ' ' + hit1 + ' ' + mrr + ' ' + ndcg + ' ' + r10 + ' ' + p95);
"
  fi
done

# Also show champion baseline for reference
echo ""
echo "Champion baseline (bge-small-en-v1.5):"
echo "  Hit@1=29.4% MRR=0.412 nDCG@10=46.0% R@10=45.1% P95=111ms"
echo ""
echo "Results saved to: $RESULTS_DIR/"