//! Implementation of the similarity index backed by KVStore.

use std::sync::Arc;

use async_trait::async_trait;
use bloomfilter::Bloom;
use parking_lot::RwLock;
use stemedb_core::types::Hash;
use tracing::{debug, instrument, warn};

use super::model::{LshBucket, MinHashSignature, SimilarityCheckResult, SimilarityIndexConfig};
use super::traits::SimilarityIndex;
use crate::error::{Result, StorageError};
use crate::key_codec;
use crate::traits::KVStore;

/// Universal hash function coefficients for MinHash.
/// Each pair (a, b) defines h(x) = (a*x + b) mod p, where p = 2^61 - 1 (Mersenne prime).
struct HashCoefficients {
    a: Vec<u64>,
    b: Vec<u64>,
}

impl HashCoefficients {
    /// Generate deterministic coefficients using BLAKE3-seeded random values.
    fn new(k: usize) -> Self {
        let mut a = Vec::with_capacity(k);
        let mut b = Vec::with_capacity(k);

        // Use BLAKE3 to deterministically generate coefficients
        for i in 0..k {
            let hash_a = blake3::hash(format!("minhash_a_{}", i).as_bytes());
            let hash_b = blake3::hash(format!("minhash_b_{}", i).as_bytes());

            // Take first 8 bytes as u64
            let mut bytes_a = [0u8; 8];
            let mut bytes_b = [0u8; 8];
            bytes_a.copy_from_slice(&hash_a.as_bytes()[..8]);
            bytes_b.copy_from_slice(&hash_b.as_bytes()[..8]);

            a.push(u64::from_le_bytes(bytes_a));
            b.push(u64::from_le_bytes(bytes_b));
        }

        Self { a, b }
    }
}

/// Mersenne prime 2^61 - 1 for universal hashing.
const MERSENNE_PRIME: u64 = (1 << 61) - 1;

/// Universal hash function: h(x) = (a*x + b) mod p
#[inline]
fn universal_hash(x: u64, a: u64, b: u64) -> u64 {
    // Use 128-bit multiplication to avoid overflow
    let ax = (a as u128) * (x as u128);
    let axb = ax + (b as u128);
    (axb % (MERSENNE_PRIME as u128)) as u64
}

/// Generic implementation of SimilarityIndex backed by any KVStore.
pub struct GenericSimilarityIndex<S> {
    store: Arc<S>,
    config: SimilarityIndexConfig,
    bloom: RwLock<Bloom<Hash>>,
    coefficients: HashCoefficients,
}

impl<S: KVStore> GenericSimilarityIndex<S> {
    /// Create a new similarity index backed by the given KV store.
    pub fn new(store: Arc<S>, config: SimilarityIndexConfig) -> Self {
        let bloom = Bloom::new_for_fp_rate(config.bloom_expected_items, config.bloom_fp_rate);
        let coefficients = HashCoefficients::new(config.minhash_k);

        Self { store, config, bloom: RwLock::new(bloom), coefficients }
    }

    /// Create a new similarity index with default configuration.
    pub fn with_defaults(store: Arc<S>) -> Self {
        Self::new(store, SimilarityIndexConfig::default())
    }

    /// Compute content hash from subject + predicate.
    fn content_hash(subject: &str, predicate: &str) -> Hash {
        let mut hasher = blake3::Hasher::new();
        hasher.update(subject.as_bytes());
        hasher.update(b":");
        hasher.update(predicate.as_bytes());
        *hasher.finalize().as_bytes()
    }

    /// Generate character n-grams (shingles) from text.
    fn shingles(&self, text: &str) -> Vec<u64> {
        let chars: Vec<char> = text.chars().collect();
        if chars.len() < self.config.shingle_size {
            // For very short text, hash the whole thing
            let hash = blake3::hash(text.as_bytes());
            let mut bytes = [0u8; 8];
            bytes.copy_from_slice(&hash.as_bytes()[..8]);
            return vec![u64::from_le_bytes(bytes)];
        }

        let mut shingles = Vec::with_capacity(chars.len() - self.config.shingle_size + 1);
        for window in chars.windows(self.config.shingle_size) {
            let s: String = window.iter().collect();
            let hash = blake3::hash(s.as_bytes());
            let mut bytes = [0u8; 8];
            bytes.copy_from_slice(&hash.as_bytes()[..8]);
            shingles.push(u64::from_le_bytes(bytes));
        }

        shingles
    }

    /// Compute MinHash signature for text.
    fn compute_minhash(&self, subject: &str, predicate: &str) -> Vec<u64> {
        let text = format!("{}:{}", subject, predicate);
        let shingles = self.shingles(&text);

        if shingles.is_empty() {
            return vec![0; self.config.minhash_k];
        }

        let mut signature = vec![u64::MAX; self.config.minhash_k];

        for shingle in shingles {
            for (sig_slot, (a, b)) in
                signature.iter_mut().zip(self.coefficients.a.iter().zip(self.coefficients.b.iter()))
            {
                let h = universal_hash(shingle, *a, *b);
                if h < *sig_slot {
                    *sig_slot = h;
                }
            }
        }

        signature
    }

    /// Compute LSH bucket hash for a band (segment of the MinHash signature).
    fn lsh_bucket_hash(&self, signature: &[u64], band: u8) -> String {
        let start = (band as usize) * self.config.lsh_rows_per_band;
        let end = start + self.config.lsh_rows_per_band;

        // Ensure we don't go out of bounds
        let end = end.min(signature.len());
        if start >= signature.len() {
            return "empty".to_string();
        }

        let band_signature = &signature[start..end];

        // Hash the band segment
        let mut hasher = blake3::Hasher::new();
        for &val in band_signature {
            hasher.update(&val.to_le_bytes());
        }

        // Return first 16 hex chars as bucket identifier
        let hash = hasher.finalize();
        hex::encode(&hash.as_bytes()[..8])
    }

    /// Get or create an LSH bucket.
    async fn get_or_create_bucket(&self, band: u8, bucket_hash: &str) -> Result<LshBucket> {
        let key = key_codec::lsh_bucket_key(band, bucket_hash);
        match self.store.get(&key).await? {
            Some(data) => stemedb_core::serde::deserialize(&data)
                .map_err(|e| StorageError::Serialization(e.to_string())),
            None => Ok(LshBucket::new()),
        }
    }

    /// Save an LSH bucket.
    async fn save_bucket(&self, band: u8, bucket_hash: &str, bucket: &LshBucket) -> Result<()> {
        let key = key_codec::lsh_bucket_key(band, bucket_hash);
        let data = stemedb_core::serde::serialize(bucket)
            .map_err(|e| StorageError::Serialization(e.to_string()))?;
        self.store.put(&key, &data).await
    }

    /// Find candidate duplicates via LSH.
    async fn find_candidates(&self, signature: &[u64]) -> Result<Vec<Hash>> {
        let mut candidates = Vec::new();

        for band in 0..self.config.lsh_bands {
            let bucket_hash = self.lsh_bucket_hash(signature, band);
            let bucket = self.get_or_create_bucket(band, &bucket_hash).await?;
            candidates.extend(bucket.members.iter().copied());
        }

        // Deduplicate candidates
        candidates.sort();
        candidates.dedup();

        Ok(candidates)
    }
}

#[async_trait]
impl<S: KVStore + 'static> SimilarityIndex for GenericSimilarityIndex<S> {
    #[instrument(skip(self), fields(subject_len = subject.len(), predicate_len = predicate.len()))]
    async fn check_similarity(
        &self,
        subject: &str,
        predicate: &str,
    ) -> Result<SimilarityCheckResult> {
        let content_hash = Self::content_hash(subject, predicate);

        // Fast Bloom filter check
        {
            let bloom = self.bloom.read();
            if !bloom.check(&content_hash) {
                debug!("Bloom filter: definitely not present");
                return Ok(SimilarityCheckResult::no_duplicate());
            }
        }

        debug!("Bloom filter: possibly present, checking MinHash");

        // Compute MinHash signature
        let signature = self.compute_minhash(subject, predicate);

        // Find candidates via LSH
        let candidates = self.find_candidates(&signature).await?;

        if candidates.is_empty() {
            debug!("No LSH candidates found");
            return Ok(SimilarityCheckResult::no_duplicate());
        }

        debug!(candidates_count = candidates.len(), "Found LSH candidates");

        // Compare with candidates
        let mut similar_entries = Vec::new();
        let mut max_similarity: f32 = 0.0;
        let mut found_exact = false;

        for candidate_hash in candidates {
            // Check for exact duplicate (same content hash already in index)
            if candidate_hash == content_hash {
                // Exact duplicate - check if we have the signature stored
                // If so, this is the same content being re-submitted
                if self.get_signature(&content_hash).await?.is_some() {
                    found_exact = true;
                    similar_entries.push(content_hash);
                    max_similarity = 1.0; // Exact match
                }
                continue;
            }

            // Get candidate signature
            if let Some(candidate_sig) = self.get_signature(&candidate_hash).await? {
                // Create temporary signature for comparison
                let temp_sig = MinHashSignature::new(
                    content_hash,
                    subject.to_string(),
                    predicate.to_string(),
                    signature.clone(),
                    0,
                );

                let similarity = temp_sig.estimate_similarity(&candidate_sig);

                if similarity >= self.config.similarity_threshold {
                    similar_entries.push(candidate_hash);
                    if similarity > max_similarity {
                        max_similarity = similarity;
                    }
                }
            }
        }

        if similar_entries.is_empty() {
            debug!("No duplicates above threshold");
            Ok(SimilarityCheckResult::no_duplicate())
        } else {
            debug!(
                similar_count = similar_entries.len(),
                max_similarity = max_similarity,
                exact_duplicate = found_exact,
                "Found duplicates"
            );
            Ok(SimilarityCheckResult::duplicate(similar_entries, max_similarity))
        }
    }

    #[instrument(skip(self), fields(subject_len = subject.len(), predicate_len = predicate.len()))]
    async fn add(&self, subject: &str, predicate: &str, timestamp: u64) -> Result<Hash> {
        let content_hash = Self::content_hash(subject, predicate);
        let hash_hex = hex::encode(content_hash);

        // Compute MinHash signature
        let signature = self.compute_minhash(subject, predicate);

        // Add to Bloom filter
        {
            let mut bloom = self.bloom.write();
            bloom.set(&content_hash);
        }

        // Store MinHash signature
        let minhash_sig = MinHashSignature::new(
            content_hash,
            subject.to_string(),
            predicate.to_string(),
            signature.clone(),
            timestamp,
        );

        let key = key_codec::minhash_key(&hash_hex);
        let data = stemedb_core::serde::serialize(&minhash_sig)
            .map_err(|e| StorageError::Serialization(e.to_string()))?;
        self.store.put(&key, &data).await?;

        // Add to LSH buckets
        for band in 0..self.config.lsh_bands {
            let bucket_hash = self.lsh_bucket_hash(&signature, band);
            let mut bucket = self.get_or_create_bucket(band, &bucket_hash).await?;
            bucket.add(content_hash);
            self.save_bucket(band, &bucket_hash, &bucket).await?;
        }

        debug!(hash = %hash_hex, "Added to similarity index");

        Ok(content_hash)
    }

    fn contains_fast(&self, content_hash: &Hash) -> bool {
        let bloom = self.bloom.read();
        bloom.check(content_hash)
    }

    #[instrument(skip(self))]
    async fn get_signature(&self, content_hash: &Hash) -> Result<Option<MinHashSignature>> {
        let hash_hex = hex::encode(content_hash);
        let key = key_codec::minhash_key(&hash_hex);

        match self.store.get(&key).await? {
            Some(data) => {
                let sig: MinHashSignature = stemedb_core::serde::deserialize(&data)
                    .map_err(|e| StorageError::Serialization(e.to_string()))?;
                Ok(Some(sig))
            }
            None => Ok(None),
        }
    }

    #[instrument(skip(self))]
    async fn len(&self) -> Result<usize> {
        let entries = self.store.scan_prefix(&key_codec::minhash_scan_prefix()).await?;
        Ok(entries.len())
    }

    #[instrument(skip(self))]
    async fn rebuild_bloom_filter(&self) -> Result<usize> {
        let entries = self.store.scan_prefix(&key_codec::minhash_scan_prefix()).await?;
        let mut count = 0;

        // Create a new Bloom filter
        let mut new_bloom =
            Bloom::new_for_fp_rate(self.config.bloom_expected_items, self.config.bloom_fp_rate);

        for (_key, data) in entries {
            match stemedb_core::serde::deserialize::<MinHashSignature>(&data) {
                Ok(sig) => {
                    new_bloom.set(&sig.content_hash);
                    count += 1;
                }
                Err(e) => {
                    warn!(error = %e, "Skipping malformed MinHash signature during rebuild");
                }
            }
        }

        // Replace the Bloom filter
        {
            let mut bloom = self.bloom.write();
            *bloom = new_bloom;
        }

        debug!(count = count, "Rebuilt Bloom filter");

        Ok(count)
    }
}

#[cfg(test)]
#[path = "tests.rs"]
mod tests;
