//! Numerical-equivalence integration test for `lunaris-embed-native`.
//!
//! Gated behind the `embedder-it` feature so it doesn't run in the default
//! `cargo test` matrix (the test needs ~623 MB of weights downloaded
//! locally — see `scripts/spike-generate-reference-embeddings.py`).
//!
//! ## Required env vars
//!
//! The test reads the on-disk path to each model artifact from these
//! environment variables. If any is missing, the test is SKIPPED (prints
//! a notice; does not fail the suite) — same convention as Lunaris's other
//! `embedder-it` tests.
//!
//! - `GRANITE_R2_WEIGHTS_PATH`   — absolute path to `model.safetensors`
//! - `GRANITE_R2_TOKENIZER_PATH` — absolute path to `tokenizer.json`
//! - `GRANITE_R2_CONFIG_PATH`    — absolute path to `config.json`
//!
//! ## Drift gates
//!
//! Asserted against the 100-prompt reference fixture
//! (`tests/fixtures/reference_embeddings.json`, generated by
//! `scripts/spike-generate-reference-embeddings.py` from HF
//! `sentence-transformers` on CPU fp32):
//!
//! - mean cosine drift ≤ 0.005 (0.5%)
//! - max  cosine drift ≤ 0.02  (2.0%)
//!
//! "Cosine drift" is defined as `1 - cosine_similarity(native, reference)`.
//! Both rows are unit-norm (the reference uses
//! `normalize_embeddings=True`; `NativeEmbedder` L2-normalizes too) so
//! cosine similarity is simply the dot product.

#![cfg(feature = "embedder-it")]

use std::path::PathBuf;

use candle_core::Device;
use lunaris_core::Embedder;
use lunaris_embed_native::{GRANITE_R2_DIM, NativeEmbedder, NativeEmbedderOpts};
use serde::Deserialize;

#[derive(Debug, Deserialize)]
struct Fixture {
    schema_version: u32,
    model_id: String,
    revision: String,
    dim: usize,
    count: usize,
    rows: Vec<FixtureRow>,
}

#[derive(Debug, Deserialize)]
struct FixtureRow {
    prompt: String,
    embedding: Vec<f64>,
}

fn env_path(name: &str) -> Option<PathBuf> {
    std::env::var_os(name).map(PathBuf::from)
}

fn load_fixture() -> Fixture {
    let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
        .join("tests")
        .join("fixtures")
        .join("reference_embeddings.json");
    let bytes =
        std::fs::read(&path).unwrap_or_else(|e| panic!("read fixture {}: {e}", path.display()));
    let fx: Fixture = serde_json::from_slice(&bytes).expect("parse fixture json");
    // why: v2 added an NFC-normalize pre-pass on both Python (reference)
    // and Rust (tokenizer.rs::encode_batch) sides so callers passing NFD
    // text get the same embedding as NFC. See P1-1 finding in
    // `.planning/phases/N-01-step-1-modernbert-fp16/P1-VERIFICATION-RESULT.md`.
    assert_eq!(fx.schema_version, 2, "fixture schema version mismatch");
    assert_eq!(fx.dim, GRANITE_R2_DIM, "fixture dim must match crate constant");
    assert_eq!(fx.rows.len(), fx.count, "fixture row count consistency");
    assert_eq!(fx.rows.len(), 100, "panel size must be exactly 100");
    fx
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn native_embedder_matches_sentence_transformers_within_drift_gate() {
    // Skip if any required env var is missing — same convention as the
    // existing lunaris-embed embedder-it tests.
    let weights = match env_path("GRANITE_R2_WEIGHTS_PATH") {
        Some(p) => p,
        None => {
            eprintln!(
                "[skip] GRANITE_R2_WEIGHTS_PATH unset; download the model with\n  \
                python3 scripts/spike-generate-reference-embeddings.py\n  \
                and export the three GRANITE_R2_*_PATH env vars."
            );
            return;
        }
    };
    let tokenizer = match env_path("GRANITE_R2_TOKENIZER_PATH") {
        Some(p) => p,
        None => {
            eprintln!("[skip] GRANITE_R2_TOKENIZER_PATH unset");
            return;
        }
    };
    let config = match env_path("GRANITE_R2_CONFIG_PATH") {
        Some(p) => p,
        None => {
            eprintln!("[skip] GRANITE_R2_CONFIG_PATH unset");
            return;
        }
    };

    let fx = load_fixture();
    eprintln!("[it] reference: {} @ {} ({} rows)", fx.model_id, fx.revision, fx.rows.len());

    let opts = NativeEmbedderOpts {
        weights_path: weights,
        tokenizer_path: tokenizer,
        config_path: config,
        device: Device::Cpu,
    };

    let embedder = NativeEmbedder::open(opts).expect("NativeEmbedder::open");
    assert_eq!(embedder.dim(), GRANITE_R2_DIM);

    // Encode the panel through NativeEmbedder. Batch size matches the
    // reference script (8) so any drift induced by batch-internal padding
    // behaves identically across the two paths.
    let prompts: Vec<&str> = fx.rows.iter().map(|r| r.prompt.as_str()).collect();
    let mut native_vecs: Vec<Vec<f32>> = Vec::with_capacity(prompts.len());
    for chunk in prompts.chunks(8) {
        let batch = embedder.embed_batch(chunk).await.expect("embed_batch");
        native_vecs.extend(batch);
    }
    assert_eq!(native_vecs.len(), fx.rows.len());

    // Per-row cosine drift. Both sides are unit-norm so cosine = dot.
    let mut drifts: Vec<f64> = Vec::with_capacity(fx.rows.len());
    let mut max_drift = 0.0_f64;
    let mut max_drift_idx = 0usize;
    for (idx, (row, native)) in fx.rows.iter().zip(native_vecs.iter()).enumerate() {
        assert_eq!(native.len(), GRANITE_R2_DIM);
        assert_eq!(row.embedding.len(), GRANITE_R2_DIM);
        let dot: f64 =
            native.iter().zip(row.embedding.iter()).map(|(a, b)| (*a as f64) * (*b)).sum();
        // Clamp to [-1, 1] to defend against accumulated f64 error pushing
        // the dot a hair over 1.0.
        let cos = dot.clamp(-1.0, 1.0);
        let drift = 1.0 - cos;
        if drift > max_drift {
            max_drift = drift;
            max_drift_idx = idx;
        }
        drifts.push(drift);
    }
    let mean: f64 = drifts.iter().sum::<f64>() / (drifts.len() as f64);

    eprintln!(
        "[it] cosine drift — mean = {:.9}, max = {:.9} (prompt #{}: {:?})",
        mean,
        max_drift,
        max_drift_idx,
        fx.rows[max_drift_idx].prompt.chars().take(60).collect::<String>()
    );

    // Discriminator: confirm the test is non-vacuous by measuring the
    // *cross-prompt* dot product (native[0] vs reference[1]). A genuine
    // model produces semantically distinct vectors for distinct prompts —
    // this dot should be in [-0.5, 0.95], never ~1.0. If it IS ~1.0 the
    // embeddings are degenerate (constant) and the drift result is
    // meaningless.
    let cross_dot: f64 = native_vecs[0]
        .iter()
        .zip(fx.rows[1].embedding.iter())
        .map(|(a, b)| (*a as f64) * (*b))
        .sum();
    eprintln!("[it] non-vacuity check — cross-prompt dot (native[0] vs ref[1]) = {cross_dot:.6}");
    eprintln!(
        "[it] sample values — native[0][0] = {:.9}, ref[0][0] = {:.9}",
        native_vecs[0][0], fx.rows[0].embedding[0]
    );
    assert!(
        cross_dot.abs() < 0.97,
        "cross-prompt dot product {cross_dot:.6} suspiciously close to 1.0 — \
         embeddings may be degenerate / constant; drift gate would be vacuous"
    );

    assert!(mean <= 0.005, "mean cosine drift {mean:.6} exceeds 0.005 gate");
    assert!(
        max_drift <= 0.02,
        "max cosine drift {max_drift:.6} exceeds 0.02 gate \
         (worst prompt #{max_drift_idx})"
    );
}
