//! Property-based tests for ID set compression.
//!
//! These tests verify mathematical invariants that must hold for all inputs,
//! using proptest to generate random test cases.

use cnk::{DeltaVarintCompressor, IdSetCompressor};
use proptest::prelude::*;

/// Generate a sorted, unique set of IDs within a universe.
fn sorted_unique_ids(max_len: usize, universe_size: u32) -> impl Strategy<Value = (Vec<u32>, u32)> {
    // Generate random subset by sampling without replacement
    (1..=max_len).prop_flat_map(move |len| {
        // Universe must be at least as large as the set
        let min_universe = len as u32;
        let actual_universe = universe_size.max(min_universe);

        proptest::collection::btree_set(0..actual_universe, len).prop_map(move |set| {
            let ids: Vec<u32> = set.into_iter().collect();
            (ids, actual_universe)
        })
    })
}

/// Generate sparse IDs (large gaps, typical of inverted indexes).
fn sparse_ids(max_len: usize) -> impl Strategy<Value = (Vec<u32>, u32)> {
    (1..=max_len).prop_flat_map(move |len| {
        // Large universe, small set
        let universe = 1_000_000u32;
        proptest::collection::btree_set(0..universe, len).prop_map(move |set| {
            let ids: Vec<u32> = set.into_iter().collect();
            (ids, universe)
        })
    })
}

/// Generate dense IDs (small gaps, typical of HNSW neighbor lists).
fn dense_ids(max_len: usize) -> impl Strategy<Value = (Vec<u32>, u32)> {
    // Generate consecutive IDs starting from a random position
    (0..10000u32, 1..=max_len).prop_map(move |(start, len)| {
        let ids: Vec<u32> = (start..start + len as u32).collect();
        let universe = start + len as u32 + 1000;
        (ids, universe)
    })
}

proptest! {
    #![proptest_config(ProptestConfig::with_cases(500))]

    // =======================================================================
    // ROUNDTRIP INVARIANT: decompress(compress(x)) == x
    // =======================================================================

    #[test]
    fn roundtrip_random_sets((ids, universe) in sorted_unique_ids(100, 10000)) {
        let compressor = DeltaVarintCompressor::new();

        let compressed = compressor.compress_set(&ids, universe)
            .expect("compression should succeed for valid input");
        let decompressed = compressor.decompress_set(&compressed, universe)
            .expect("decompression should succeed for valid compressed data");

        prop_assert_eq!(ids, decompressed, "roundtrip must preserve data");
    }

    #[test]
    fn roundtrip_sparse_sets((ids, universe) in sparse_ids(50)) {
        let compressor = DeltaVarintCompressor::new();

        let compressed = compressor.compress_set(&ids, universe)?;
        let decompressed = compressor.decompress_set(&compressed, universe)?;

        prop_assert_eq!(ids, decompressed);
    }

    #[test]
    fn roundtrip_dense_sets((ids, universe) in dense_ids(100)) {
        let compressor = DeltaVarintCompressor::new();

        let compressed = compressor.compress_set(&ids, universe)?;
        let decompressed = compressor.decompress_set(&compressed, universe)?;

        prop_assert_eq!(ids, decompressed);
    }

    // =======================================================================
    // SIZE BOUNDS: compressed size should be reasonable
    // =======================================================================

    #[test]
    fn compression_reduces_size_for_large_sets((ids, universe) in sorted_unique_ids(100, 100000)) {
        prop_assume!(!ids.is_empty());

        let compressor = DeltaVarintCompressor::new();
        let compressed = compressor.compress_set(&ids, universe)?;

        let uncompressed_size = ids.len() * 4; // 4 bytes per u32
        let compressed_size = compressed.len();

        // For non-trivial sets, compressed should be smaller than raw
        // (varint on deltas is almost always better than raw u32)
        if ids.len() > 10 {
            prop_assert!(
                compressed_size <= uncompressed_size,
                "compressed {} should be <= uncompressed {} for {} IDs",
                compressed_size, uncompressed_size, ids.len()
            );
        }
    }

    #[test]
    fn consecutive_ids_compress_well((start, len) in (0u32..10000, 10usize..200)) {
        let ids: Vec<u32> = (start..start + len as u32).collect();
        let universe = start + len as u32 + 1000;

        let compressor = DeltaVarintCompressor::new();
        let compressed = compressor.compress_set(&ids, universe)?;

        // Consecutive IDs have delta=1, which is 1 byte per ID in varint
        // Plus header (count + first ID), so expect ~1-2 bytes per ID
        let bytes_per_id = compressed.len() as f64 / ids.len() as f64;
        prop_assert!(
            bytes_per_id < 2.5,
            "consecutive IDs should compress to ~1-2 bytes/ID, got {}",
            bytes_per_id
        );
    }

    // =======================================================================
    // EDGE CASES
    // =======================================================================

    #[test]
    fn empty_set_roundtrip(universe in 1u32..1000000) {
        let compressor = DeltaVarintCompressor::new();
        let ids: Vec<u32> = vec![];

        let compressed = compressor.compress_set(&ids, universe)?;
        let decompressed = compressor.decompress_set(&compressed, universe)?;

        prop_assert!(compressed.is_empty(), "empty set should compress to empty");
        prop_assert!(decompressed.is_empty(), "empty should decompress to empty");
    }

    #[test]
    fn single_id_roundtrip(id in 0u32..100000) {
        let universe = id + 1;
        let compressor = DeltaVarintCompressor::new();
        let ids = vec![id];

        let compressed = compressor.compress_set(&ids, universe)?;
        let decompressed = compressor.decompress_set(&compressed, universe)?;

        prop_assert_eq!(ids, decompressed);
    }

    #[test]
    fn two_ids_roundtrip((a, gap) in (0u32..10000, 1u32..10000)) {
        let b = a.saturating_add(gap);
        let universe = b + 1;
        let compressor = DeltaVarintCompressor::new();
        let ids = vec![a, b];

        let compressed = compressor.compress_set(&ids, universe)?;
        let decompressed = compressor.decompress_set(&compressed, universe)?;

        prop_assert_eq!(ids, decompressed);
    }

    // =======================================================================
    // ERROR CASES
    // =======================================================================

    #[test]
    fn rejects_unsorted_ids(
        (a, b) in (1u32..10000, 0u32..10000)
    ) {
        prop_assume!(a > b); // Ensure unsorted
        let compressor = DeltaVarintCompressor::new();
        let ids = vec![a, b]; // Intentionally unsorted

        let result = compressor.compress_set(&ids, 100000);
        prop_assert!(result.is_err(), "should reject unsorted IDs");
    }

    #[test]
    fn rejects_ids_exceeding_universe((ids, _universe) in sorted_unique_ids(10, 1000)) {
        prop_assume!(!ids.is_empty());
        let max_id = *ids.iter().max().unwrap();

        // Use a universe smaller than max_id
        let small_universe = max_id; // max_id is now out of bounds

        let compressor = DeltaVarintCompressor::new();
        let result = compressor.compress_set(&ids, small_universe);

        prop_assert!(result.is_err(), "should reject IDs >= universe");
    }

    // =======================================================================
    // DETERMINISM
    // =======================================================================

    #[test]
    fn compression_is_deterministic((ids, universe) in sorted_unique_ids(50, 10000)) {
        let compressor = DeltaVarintCompressor::new();

        let compressed1 = compressor.compress_set(&ids, universe)?;
        let compressed2 = compressor.compress_set(&ids, universe)?;

        prop_assert_eq!(compressed1, compressed2, "compression must be deterministic");
    }

    #[cfg(feature = "ans")]
    #[test]
    fn roc_roundtrip_random_sets((ids, universe) in sorted_unique_ids(100, 10000)) {
        let compressor = cnk::RocCompressor::new();

        let compressed = compressor
            .compress_set(&ids, universe)
            .expect("compression should succeed for valid input");
        let decompressed = compressor
            .decompress_set(&compressed, universe)
            .expect("decompression should succeed for valid compressed data");

        prop_assert_eq!(ids, decompressed, "ROC roundtrip must preserve data");
    }

    #[cfg(feature = "ans")]
    #[test]
    fn roc_roundtrip_sparse_sets((ids, universe) in sparse_ids(50)) {
        let compressor = cnk::RocCompressor::new();

        let compressed = compressor.compress_set(&ids, universe)?;
        let decompressed = compressor.decompress_set(&compressed, universe)?;

        prop_assert_eq!(ids, decompressed);
    }

    #[cfg(feature = "ans")]
    #[test]
    fn roc_compression_is_deterministic((ids, universe) in sorted_unique_ids(50, 10000)) {
        let compressor = cnk::RocCompressor::new();

        let compressed1 = compressor.compress_set(&ids, universe)?;
        let compressed2 = compressor.compress_set(&ids, universe)?;

        prop_assert_eq!(compressed1, compressed2, "ROC compression must be deterministic");
    }

    #[cfg(feature = "sbits")]
    #[test]
    fn elias_fano_roundtrip_random_sets((ids, universe) in sorted_unique_ids(100, 10000)) {
        let compressor = cnk::EliasFanoCompressor::new();

        let compressed = compressor
            .compress_set(&ids, universe)
            .expect("compression should succeed for valid input");
        let decompressed = compressor
            .decompress_set(&compressed, universe)
            .expect("decompression should succeed for valid compressed data");

        prop_assert_eq!(ids, decompressed, "Elias-Fano roundtrip must preserve data");
    }

    #[cfg(feature = "sbits")]
    #[test]
    fn partitioned_elias_fano_roundtrip_random_sets((ids, universe) in sorted_unique_ids(100, 10000)) {
        let compressor = cnk::PartitionedEliasFanoCompressor::with_block_size(32);

        let compressed = compressor
            .compress_set(&ids, universe)
            .expect("compression should succeed for valid input");
        let decompressed = compressor
            .decompress_set(&compressed, universe)
            .expect("decompression should succeed for valid compressed data");

        prop_assert_eq!(ids, decompressed, "Partitioned EF roundtrip must preserve data");
    }

    // =======================================================================
    // VARINT EDGE CASES
    // =======================================================================

    #[test]
    fn handles_large_ids(start in 0u32..u32::MAX - 1000) {
        let ids: Vec<u32> = (0..10).map(|i| start.saturating_add(i * 10)).collect();
        let universe = ids.last().unwrap() + 1;

        let compressor = DeltaVarintCompressor::new();
        let compressed = compressor.compress_set(&ids, universe)?;
        let decompressed = compressor.decompress_set(&compressed, universe)?;

        prop_assert_eq!(ids, decompressed);
    }

    #[test]
    fn handles_large_gaps((a, gap) in (0u32..1000, 1u32..u32::MAX / 2)) {
        let b = a.saturating_add(gap);
        let universe = b + 1;
        let compressor = DeltaVarintCompressor::new();
        let ids = vec![a, b];

        let compressed = compressor.compress_set(&ids, universe)?;
        let decompressed = compressor.decompress_set(&compressed, universe)?;

        prop_assert_eq!(ids, decompressed);
    }
}

proptest! {
    #![proptest_config(ProptestConfig::with_cases(200))]

    // =======================================================================
    // ADDITIONAL ROUNDTRIP: large universe, sorted+deduped input
    // =======================================================================

    #[test]
    fn prop_delta_varint_roundtrip_deduped(
        ids in prop::collection::vec(0u32..10_000, 0..500),
        universe_size in 10_001u32..100_000,
    ) {
        let compressor = DeltaVarintCompressor::new();
        let mut ids = ids;
        ids.sort_unstable();
        ids.dedup();

        let compressed = compressor.compress_set(&ids, universe_size)?;
        let decompressed = compressor.decompress_set(&compressed, universe_size)?;
        prop_assert_eq!(&ids, &decompressed);
    }

    // =======================================================================
    // COMPRESSION RATIO: delta encoding beats raw u32 for non-trivial sets
    // =======================================================================

    #[test]
    fn prop_compression_ratio(
        ids in prop::collection::vec(0u32..100_000, 100..1000),
        universe_size_offset in 1u32..100_000,
    ) {
        let compressor = DeltaVarintCompressor::new();
        let mut ids = ids;
        ids.sort_unstable();
        ids.dedup();
        if ids.is_empty() { return Ok(()); }
        let max_id = *ids.last().unwrap();
        let universe_size = max_id.saturating_add(universe_size_offset).saturating_add(1);

        let compressed = compressor.compress_set(&ids, universe_size)?;
        let raw_size = ids.len() * 4;
        // Delta encoding should always beat raw u32 storage for sorted sets
        prop_assert!(compressed.len() <= raw_size + 10,
            "compressed {} bytes > raw {} bytes + margin", compressed.len(), raw_size);
    }

    // =======================================================================
    // BITS_PER_ID CORRELATION: estimate correlates with actual compressed size
    // =======================================================================

    #[test]
    fn prop_estimate_correlates_with_actual_size(
        ids in prop::collection::vec(0u32..100_000, 20..500),
        universe_size_offset in 1u32..100_000,
    ) {
        let compressor = DeltaVarintCompressor::new();
        let mut ids = ids;
        ids.sort_unstable();
        ids.dedup();
        if ids.len() < 10 { return Ok(()); }
        let max_id = *ids.last().unwrap();
        let universe_size = max_id.saturating_add(universe_size_offset).saturating_add(1);

        let compressed = compressor.compress_set(&ids, universe_size)?;
        let actual_bytes = compressed.len();
        let estimated_bytes = compressor.estimate_size(ids.len(), universe_size);

        // The estimate uses the mean-gap model (uniform distribution assumption).
        // For clustered data, actual may be smaller (small gaps = 1-byte varints).
        // For sparse data, actual may be larger (large gaps = multi-byte varints).
        // Allow 3x in either direction -- the estimate is a planning tool, not exact.
        prop_assert!(
            actual_bytes <= estimated_bytes * 3 + 16,
            "actual {} bytes vastly exceeds 3x estimated {} bytes for {} IDs in universe {}",
            actual_bytes, estimated_bytes, ids.len(), universe_size
        );
        prop_assert!(
            estimated_bytes <= actual_bytes * 3 + 16,
            "estimated {} bytes vastly exceeds 3x actual {} bytes for {} IDs in universe {}",
            estimated_bytes, actual_bytes, ids.len(), universe_size
        );
    }
}

// =======================================================================
// STATISTICAL TESTS (not proptest, but important)
// =======================================================================

#[test]
fn compression_ratio_improves_with_density() {
    let compressor = DeltaVarintCompressor::new();
    let universe = 100_000u32;

    // Sparse set with large gaps (delta ~100)
    let sparse: Vec<u32> = (0..1000).map(|i| i * 100).collect();
    let sparse_compressed = compressor.compress_set(&sparse, universe).unwrap();
    let sparse_bytes_per_id = sparse_compressed.len() as f64 / sparse.len() as f64;

    // Dense set (consecutive, delta = 1)
    let dense: Vec<u32> = (0..1000).collect();
    let dense_compressed = compressor.compress_set(&dense, universe).unwrap();
    let dense_bytes_per_id = dense_compressed.len() as f64 / dense.len() as f64;

    // Dense should use fewer bytes per ID (delta=1 is 1 byte, delta=100 is 1 byte too in varint)
    // Actually both fit in 1 byte varint, so let's use larger gaps
    let very_sparse: Vec<u32> = (0..100).map(|i| i * 1000).collect();
    let very_sparse_compressed = compressor.compress_set(&very_sparse, universe).unwrap();
    let very_sparse_bytes_per_id = very_sparse_compressed.len() as f64 / very_sparse.len() as f64;

    // Very sparse (delta=1000, needs 2 bytes) should use more bytes per ID than dense (delta=1)
    assert!(
        very_sparse_bytes_per_id > dense_bytes_per_id,
        "very sparse {} bytes/id should exceed dense {} bytes/id",
        very_sparse_bytes_per_id,
        dense_bytes_per_id
    );

    println!("Dense: {:.2} bytes/id", dense_bytes_per_id);
    println!("Sparse (delta=100): {:.2} bytes/id", sparse_bytes_per_id);
    println!(
        "Very sparse (delta=1000): {:.2} bytes/id",
        very_sparse_bytes_per_id
    );
}

#[test]
fn estimate_size_vs_actual() {
    let compressor = DeltaVarintCompressor::new();

    for num_ids in [10, 100, 1000] {
        for universe in [10_000u32, 100_000, 1_000_000] {
            let estimate = compressor.estimate_size(num_ids, universe);

            // Estimate should be positive for non-empty sets
            assert!(estimate > 0, "estimate should be positive");

            // Estimate should be less than raw storage
            let raw_size = num_ids * 4;
            assert!(
                estimate <= raw_size,
                "estimate {} should not exceed raw {} for {} ids in universe {}",
                estimate,
                raw_size,
                num_ids,
                universe
            );

            // Compare against actual compression for uniform-ish data
            let ids: Vec<u32> = (0..num_ids as u32)
                .map(|i| i * (universe / num_ids as u32))
                .collect();
            let compressed = compressor.compress_set(&ids, universe).unwrap();
            let actual = compressed.len();

            // For uniformly-spaced IDs, estimate should be within 2x of actual
            assert!(
                estimate <= actual * 2 + 16,
                "estimate {} too high vs actual {} for {} uniform ids in universe {}",
                estimate,
                actual,
                num_ids,
                universe
            );
            assert!(
                actual <= estimate * 2 + 16,
                "actual {} too high vs estimate {} for {} uniform ids in universe {}",
                actual,
                estimate,
                num_ids,
                universe
            );
        }
    }
}

// =======================================================================
// ERROR PATH COVERAGE
// =======================================================================

#[test]
fn decompress_truncated_data() {
    let compressor = DeltaVarintCompressor::new();
    let ids = vec![1u32, 5, 10, 20, 50];
    let compressed = compressor.compress_set(&ids, 100).unwrap();

    // Truncate to just the count byte
    let truncated = &compressed[..1];
    let result = compressor.decompress_set(truncated, 100);
    assert!(result.is_err(), "truncated data should fail");
}

#[test]
fn decompress_trailing_bytes() {
    let compressor = DeltaVarintCompressor::new();
    let ids = vec![1u32, 5, 10];
    let mut compressed = compressor.compress_set(&ids, 100).unwrap();
    compressed.push(0xFF); // trailing garbage

    let result = compressor.decompress_set(&compressed, 100);
    assert!(result.is_err(), "trailing bytes should fail");
}

#[test]
fn decompress_num_ids_exceeds_universe() {
    let compressor = DeltaVarintCompressor::new();
    // Craft bytes: count=1000 (varint), but universe_size=10
    // varint 1000 = 0xe8 0x07
    let crafted = vec![0xe8, 0x07, 0x00]; // count=1000, first_id=0
    let result = compressor.decompress_set(&crafted, 10);
    assert!(result.is_err(), "num_ids exceeding universe should fail");
}

#[test]
fn compress_universe_zero_empty_set() {
    let compressor = DeltaVarintCompressor::new();
    let result = compressor.compress_set(&[], 0);
    assert!(result.is_ok(), "empty set with universe 0 should succeed");
}

#[test]
fn compress_universe_zero_nonempty() {
    let compressor = DeltaVarintCompressor::new();
    let result = compressor.compress_set(&[0], 0);
    assert!(result.is_err(), "non-empty set with universe 0 should fail");
}

#[test]
fn envelope_truncated_header() {
    let result = cnk::decompress_set_enveloped(b"CNK");
    assert!(result.is_err(), "truncated envelope should fail");
}

#[test]
fn envelope_bad_magic() {
    let result = cnk::decompress_set_enveloped(b"BADMAGIC00000000000000000000000000");
    assert!(result.is_err(), "bad magic should fail");
}

#[test]
fn auto_roundtrip_no_sbits_fallback() {
    // Test the auto path without sbits features -- should always use DeltaVarint
    let ids: Vec<u32> = (0..100).collect();
    let (choice, bytes) = cnk::compress_set_auto(&ids, 1000, cnk::ChooseConfig::default()).unwrap();
    let back = cnk::decompress_set_auto(choice, &bytes, 1000).unwrap();
    assert_eq!(ids, back);
}
