//! CRDT wrapper for assertion storage implementing G-Set semantics.
//!
//! Assertions naturally form a G-Set (Grow-only Set):
//! - Assertions are append-only (never deleted)
//! - Content-addressed by BLAKE3 hash (idempotent inserts)
//!
//! This wrapper adds explicit merge operations for replication.

use crate::error::{Result, StorageError};
use crate::key_codec;
use crate::traits::KVStore;
use async_trait::async_trait;
use rkyv::{Archive, Deserialize, Serialize};
use std::sync::Arc;
use stemedb_core::types::Hash;
use tracing::{debug, instrument, warn};

use super::traits::CrdtMerge;

/// G-Set state for assertions under a subject.
///
/// This is a set of assertion hashes - the actual assertion data
/// is content-addressed and can be fetched separately.
#[derive(Archive, Deserialize, Serialize, Debug, Clone, PartialEq)]
#[archive(check_bytes)]
pub struct AssertionSetState {
    /// The subject this state covers.
    pub subject: String,
    /// Set of assertion hashes present on this node.
    pub assertion_hashes: Vec<Hash>,
    /// Source node ID.
    pub source_node: [u8; 16],
}

impl AssertionSetState {
    /// Creates a new assertion set state.
    pub fn new(subject: String, assertion_hashes: Vec<Hash>, source_node: [u8; 16]) -> Self {
        Self { subject, assertion_hashes, source_node }
    }

    /// Merges two assertion set states using set union.
    ///
    /// G-Set merge is simply the union of both sets.
    pub fn merge(&self, other: &Self) -> Self {
        debug_assert_eq!(self.subject, other.subject, "Cannot merge states for different subjects");

        // Union of both hash sets
        let mut combined: Vec<Hash> = self.assertion_hashes.clone();
        for hash in &other.assertion_hashes {
            if !combined.contains(hash) {
                combined.push(*hash);
            }
        }

        Self {
            subject: self.subject.clone(),
            assertion_hashes: combined,
            source_node: self.source_node,
        }
    }

    /// Returns the number of assertions in this state.
    pub fn len(&self) -> usize {
        self.assertion_hashes.len()
    }

    /// Returns true if the state is empty.
    pub fn is_empty(&self) -> bool {
        self.assertion_hashes.is_empty()
    }

    /// Checks if an assertion hash is in this state.
    pub fn contains(&self, hash: &Hash) -> bool {
        self.assertion_hashes.contains(hash)
    }
}

/// Assertion data that may need to be transferred during sync.
///
/// When merging G-Sets, we first compare hashes. If the remote has
/// hashes we don't have, we request the full assertion data.
#[derive(Archive, Deserialize, Serialize, Debug, Clone)]
#[archive(check_bytes)]
pub struct AssertionTransfer {
    /// The assertion hash (for verification).
    pub hash: Hash,
    /// The raw serialized assertion bytes.
    pub data: Vec<u8>,
}

/// CRDT wrapper for assertion storage with G-Set merge semantics.
///
/// Wraps a KVStore and adds merge operations for distributed replication.
/// Assertions are content-addressed by their BLAKE3 hash, making inserts
/// naturally idempotent.
///
/// # Merge Semantics
///
/// Assertion sets use G-Set (union) semantics:
/// - `merge(A, B)` = `A ∪ B` (set union)
/// - Missing assertions are requested and stored locally
///
/// # Example
///
/// ```ignore
/// use stemedb_storage::crdt::CrdtAssertionStore;
/// use std::sync::Arc;
///
/// let crdt = CrdtAssertionStore::new(Arc::new(store), node_id);
///
/// // Get set of assertion hashes for a subject
/// let state = crdt.get_state("Tesla_Inc").await?;
///
/// // Compare with remote state to find missing assertions
/// let missing = crdt.find_missing("Tesla_Inc", &remote_state).await?;
///
/// // Merge remote assertions (with their data)
/// crdt.merge_with_data("Tesla_Inc", &remote_assertions).await?;
/// ```
pub struct CrdtAssertionStore<S: KVStore> {
    store: Arc<S>,
    node_id: [u8; 16],
}

impl<S: KVStore + 'static> CrdtAssertionStore<S> {
    /// Creates a new CRDT assertion store with the given node ID.
    pub fn new(store: Arc<S>, node_id: [u8; 16]) -> Self {
        Self { store, node_id }
    }

    /// Gets an assertion by its hash.
    #[instrument(skip(self), fields(hash = %hex::encode(hash)))]
    pub async fn get_assertion(&self, subject: &str, hash: &Hash) -> Result<Option<Vec<u8>>> {
        let hash_hex = hex::encode(hash);
        let key = key_codec::assertion_key(subject, &hash_hex);
        self.store.get(&key).await
    }

    /// Puts an assertion (content-addressed, idempotent).
    ///
    /// The hash is computed from the data, so duplicate puts are safe.
    #[instrument(skip(self, data), fields(data_len = data.len()))]
    pub async fn put_assertion(&self, subject: &str, data: &[u8]) -> Result<Hash> {
        let hash_bytes = blake3::hash(data);
        let hash: Hash = *hash_bytes.as_bytes();
        let hash_hex = hex::encode(hash);

        let key = key_codec::assertion_key(subject, &hash_hex);
        self.store.put(&key, data).await?;

        debug!(hash = %hash_hex, "Stored assertion");
        Ok(hash)
    }

    /// Checks if an assertion exists locally.
    #[instrument(skip(self))]
    pub async fn has_assertion(&self, subject: &str, hash: &Hash) -> Result<bool> {
        let hash_hex = hex::encode(hash);
        let key = key_codec::assertion_key(subject, &hash_hex);
        Ok(self.store.get(&key).await?.is_some())
    }

    /// Finds assertion hashes present in remote state but missing locally.
    ///
    /// Returns hashes that need to be fetched from the remote node.
    #[instrument(skip(self, remote), fields(remote_count = remote.assertion_hashes.len()))]
    pub async fn find_missing(
        &self,
        subject: &str,
        remote: &AssertionSetState,
    ) -> Result<Vec<Hash>> {
        if remote.subject != subject {
            return Err(StorageError::InputValidation("Subject mismatch".to_string()));
        }

        let mut missing = Vec::new();
        for hash in &remote.assertion_hashes {
            if !self.has_assertion(subject, hash).await? {
                missing.push(*hash);
            }
        }

        debug!(missing_count = missing.len(), "Found missing assertions");
        Ok(missing)
    }

    /// Merges assertion data received from a remote node.
    ///
    /// Each assertion is verified by computing its hash and comparing
    /// to the expected hash before storing.
    #[instrument(skip(self, assertions), fields(count = assertions.len()))]
    pub async fn merge_with_data(
        &self,
        subject: &str,
        assertions: &[AssertionTransfer],
    ) -> Result<usize> {
        let mut merged_count = 0;

        for transfer in assertions {
            // Verify hash
            let computed_hash = blake3::hash(&transfer.data);
            if computed_hash.as_bytes() != &transfer.hash {
                warn!(
                    expected = %hex::encode(transfer.hash),
                    computed = %hex::encode(computed_hash.as_bytes()),
                    "Hash mismatch in assertion transfer, skipping"
                );
                continue;
            }

            // Store if not already present
            if !self.has_assertion(subject, &transfer.hash).await? {
                self.put_assertion(subject, &transfer.data).await?;
                merged_count += 1;
            }
        }

        debug!(merged_count, "Merged assertion data from remote");
        Ok(merged_count)
    }
}

#[async_trait]
impl<S: KVStore + 'static> CrdtMerge for CrdtAssertionStore<S> {
    type State = AssertionSetState;

    #[instrument(skip(self))]
    async fn get_state(&self, subject: &str) -> Result<Self::State> {
        // Scan all assertion keys for this subject
        let prefix = key_codec::assertion_prefix(subject);
        let entries = self.store.scan_prefix(&prefix).await?;

        let mut hashes = Vec::with_capacity(entries.len());

        for (key, _) in entries {
            // Extract hash from key
            // Key format: {subject}\x00H:{hash_hex}
            let key_str = String::from_utf8_lossy(&key);
            if let Some(hash_hex) = key_str.split(':').next_back() {
                if let Ok(hash_bytes) = hex::decode(hash_hex) {
                    if hash_bytes.len() == 32 {
                        let hash: Hash = hash_bytes.try_into().map_err(|_| {
                            StorageError::Serialization("Invalid hash bytes".to_string())
                        })?;
                        hashes.push(hash);
                    }
                }
            }
        }

        Ok(AssertionSetState::new(subject.to_string(), hashes, self.node_id))
    }

    #[instrument(skip(self, remote), fields(subject = %remote.subject, hash_count = remote.assertion_hashes.len()))]
    async fn merge(&self, subject: &str, remote: &Self::State) -> Result<()> {
        if remote.subject != subject {
            warn!(
                expected = subject,
                actual = %remote.subject,
                "Subject mismatch in merge"
            );
            return Err(StorageError::InputValidation("Subject mismatch in merge".to_string()));
        }

        // G-Set merge: just need to ensure all hashes exist
        // The actual data transfer is handled separately via merge_with_data
        let missing = self.find_missing(subject, remote).await?;

        if !missing.is_empty() {
            debug!(
                missing_count = missing.len(),
                "Merge found missing assertions - data transfer required"
            );
            // Note: Caller is responsible for fetching and merging the actual data
            // using merge_with_data(). This method only identifies what's missing.
        }

        Ok(())
    }

    fn node_id(&self) -> [u8; 16] {
        self.node_id
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::HybridStore;
    use tempfile::tempdir;

    async fn create_test_store() -> Arc<HybridStore> {
        let dir = tempdir().expect("Failed to create temp dir");
        Arc::new(HybridStore::open(dir.path()).expect("Failed to open store"))
    }

    #[tokio::test]
    async fn test_assertion_set_state_merge() {
        let hash1 = [1u8; 32];
        let hash2 = [2u8; 32];
        let hash3 = [3u8; 32];
        let node1 = [1u8; 16];
        let node2 = [2u8; 16];

        let state1 = AssertionSetState::new("test".to_string(), vec![hash1, hash2], node1);
        let state2 = AssertionSetState::new("test".to_string(), vec![hash2, hash3], node2);

        let merged = state1.merge(&state2);

        assert_eq!(merged.assertion_hashes.len(), 3);
        assert!(merged.contains(&hash1));
        assert!(merged.contains(&hash2));
        assert!(merged.contains(&hash3));
    }

    #[tokio::test]
    async fn test_put_and_get_assertion() {
        let store = create_test_store().await;
        let crdt = CrdtAssertionStore::new(store, [1u8; 16]);

        let data = b"test assertion data";
        let hash = crdt.put_assertion("test", data).await.expect("put");

        let retrieved = crdt.get_assertion("test", &hash).await.expect("get");
        assert!(retrieved.is_some());
        assert_eq!(retrieved.expect("should exist"), data.to_vec());
    }

    #[tokio::test]
    async fn test_put_is_idempotent() {
        let store = create_test_store().await;
        let crdt = CrdtAssertionStore::new(store, [1u8; 16]);

        let data = b"test assertion data";

        let hash1 = crdt.put_assertion("test", data).await.expect("put1");
        let hash2 = crdt.put_assertion("test", data).await.expect("put2");

        // Same data = same hash
        assert_eq!(hash1, hash2);

        // Only one entry
        let state = crdt.get_state("test").await.expect("state");
        assert_eq!(state.len(), 1);
    }

    #[tokio::test]
    async fn test_find_missing() {
        let store = create_test_store().await;
        let crdt = CrdtAssertionStore::new(store, [1u8; 16]);

        // Put one assertion locally
        let local_data = b"local assertion";
        let local_hash = crdt.put_assertion("test", local_data).await.expect("put");

        // Remote state has local + one more
        let remote_only_hash = [99u8; 32];
        let remote_state = AssertionSetState::new(
            "test".to_string(),
            vec![local_hash, remote_only_hash],
            [2u8; 16],
        );

        let missing = crdt.find_missing("test", &remote_state).await.expect("find");

        assert_eq!(missing.len(), 1);
        assert_eq!(missing[0], remote_only_hash);
    }

    #[tokio::test]
    async fn test_merge_with_data() {
        let store = create_test_store().await;
        let crdt = CrdtAssertionStore::new(store, [1u8; 16]);

        let data = b"transferred assertion";
        let hash_bytes = blake3::hash(data);
        let hash: Hash = *hash_bytes.as_bytes();

        let transfers = vec![AssertionTransfer { hash, data: data.to_vec() }];

        let merged = crdt.merge_with_data("test", &transfers).await.expect("merge");
        assert_eq!(merged, 1);

        // Should now exist
        let exists = crdt.has_assertion("test", &hash).await.expect("has");
        assert!(exists);

        // Merging again should be idempotent
        let merged2 = crdt.merge_with_data("test", &transfers).await.expect("merge2");
        assert_eq!(merged2, 0); // Already exists
    }

    #[tokio::test]
    async fn test_merge_with_data_rejects_bad_hash() {
        let store = create_test_store().await;
        let crdt = CrdtAssertionStore::new(store, [1u8; 16]);

        let data = b"some data";
        let wrong_hash = [0u8; 32]; // Doesn't match data

        let transfers = vec![AssertionTransfer { hash: wrong_hash, data: data.to_vec() }];

        let merged = crdt.merge_with_data("test", &transfers).await.expect("merge");
        assert_eq!(merged, 0); // Should reject due to hash mismatch
    }
}

#[cfg(test)]
mod property_tests {
    use super::*;
    use proptest::prelude::*;

    // Property: AssertionSetState merge is commutative
    proptest! {
        #[test]
        fn merge_commutative(
            hashes_a in prop::collection::vec(prop::array::uniform32(0u8..255), 0..10),
            hashes_b in prop::collection::vec(prop::array::uniform32(0u8..255), 0..10),
        ) {
            let node1 = [1u8; 16];
            let node2 = [2u8; 16];

            let state_a = AssertionSetState::new("test".to_string(), hashes_a.clone(), node1);
            let state_b = AssertionSetState::new("test".to_string(), hashes_b.clone(), node2);

            let merged_ab = state_a.merge(&state_b);
            let merged_ba = state_b.merge(&state_a);

            // Same hashes regardless of order
            let mut ab_sorted = merged_ab.assertion_hashes.clone();
            let mut ba_sorted = merged_ba.assertion_hashes.clone();
            ab_sorted.sort();
            ba_sorted.sort();

            prop_assert_eq!(ab_sorted, ba_sorted);
        }
    }

    // Property: AssertionSetState merge is associative
    proptest! {
        #[test]
        fn merge_associative(
            hashes_a in prop::collection::vec(prop::array::uniform32(0u8..255), 0..5),
            hashes_b in prop::collection::vec(prop::array::uniform32(0u8..255), 0..5),
            hashes_c in prop::collection::vec(prop::array::uniform32(0u8..255), 0..5),
        ) {
            let node1 = [1u8; 16];
            let node2 = [2u8; 16];
            let node3 = [3u8; 16];

            let state_a = AssertionSetState::new("test".to_string(), hashes_a, node1);
            let state_b = AssertionSetState::new("test".to_string(), hashes_b, node2);
            let state_c = AssertionSetState::new("test".to_string(), hashes_c, node3);

            let merged_ab_c = state_a.merge(&state_b).merge(&state_c);
            let merged_a_bc = state_a.merge(&state_b.merge(&state_c));

            let mut ab_c_sorted = merged_ab_c.assertion_hashes.clone();
            let mut a_bc_sorted = merged_a_bc.assertion_hashes.clone();
            ab_c_sorted.sort();
            a_bc_sorted.sort();

            prop_assert_eq!(ab_c_sorted, a_bc_sorted);
        }
    }

    // Property: AssertionSetState merge is idempotent
    proptest! {
        #[test]
        fn merge_idempotent(
            hashes in prop::collection::vec(prop::array::uniform32(0u8..255), 0..10),
        ) {
            let node_id = [1u8; 16];
            let state = AssertionSetState::new("test".to_string(), hashes, node_id);
            let merged = state.merge(&state);

            let mut original_sorted = state.assertion_hashes.clone();
            let mut merged_sorted = merged.assertion_hashes.clone();
            original_sorted.sort();
            merged_sorted.sort();

            prop_assert_eq!(original_sorted, merged_sorted);
        }
    }
}
