//! CRDT wrapper for VoteStore implementing G-Counter semantics.
//!
//! The vote store naturally implements G-Counter (Grow-only Counter) semantics:
//! - Vote counts only increase
//! - Aggregate weights only increase (assuming positive weights)
//!
//! This wrapper adds explicit merge operations for replication.

use crate::error::Result;
use crate::key_codec;
use crate::traits::KVStore;
use crate::vote_store::{GenericVoteStore, VoteStore};
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-Counter state for vote counts per assertion.
///
/// This state captures the vote count and aggregate weight for a single
/// assertion within a subject. It's designed for efficient delta sync.
#[derive(Archive, Deserialize, Serialize, Debug, Clone, PartialEq)]
#[archive(check_bytes)]
pub struct VoteCountState {
    /// The assertion this state applies to.
    pub assertion_hash: Hash,
    /// Total vote count from all nodes.
    pub count: u64,
    /// Aggregate weight from all nodes.
    pub weight: f32,
    /// Node that produced this state (for provenance).
    pub source_node: [u8; 16],
}

impl VoteCountState {
    /// Creates a new vote count state.
    pub fn new(assertion_hash: Hash, count: u64, weight: f32, source_node: [u8; 16]) -> Self {
        Self { assertion_hash, count, weight, source_node }
    }

    /// Merges two vote count states, taking the maximum of each field.
    ///
    /// This implements G-Counter merge semantics where counts only grow.
    pub fn merge(&self, other: &Self) -> Self {
        debug_assert_eq!(
            self.assertion_hash, other.assertion_hash,
            "Cannot merge states for different assertions"
        );

        Self {
            assertion_hash: self.assertion_hash,
            count: self.count.max(other.count),
            weight: self.weight.max(other.weight),
            source_node: self.source_node, // Keep local node ID
        }
    }
}

/// CRDT wrapper for VoteStore with G-Counter merge semantics.
///
/// Wraps a `GenericVoteStore` and adds merge operations for distributed
/// replication. The underlying atomic operations (`fetch_and_add_u64`,
/// `compare_and_swap_f32`) already provide local consistency; this wrapper
/// adds cross-node consistency via explicit merge.
///
/// # Merge Semantics
///
/// Vote counts use G-Counter (max) semantics:
/// - `merge(local, remote)` takes `max(local.count, remote.count)`
/// - This ensures counts converge to the highest observed value
///
/// # Example
///
/// ```ignore
/// use stemedb_storage::crdt::CrdtVoteStore;
/// use std::sync::Arc;
///
/// let crdt = CrdtVoteStore::new(Arc::new(store), node_id);
///
/// // Local operations work as normal
/// crdt.put_vote(&vote, "subject").await?;
///
/// // Get state to send to remote node
/// let state = crdt.get_state("subject").await?;
///
/// // On receiving remote state, merge it
/// crdt.merge("subject", &remote_state).await?;
/// ```
pub struct CrdtVoteStore<S: KVStore> {
    inner: GenericVoteStore<Arc<S>>,
    store: Arc<S>,
    node_id: [u8; 16],
}

impl<S: KVStore + 'static> CrdtVoteStore<S> {
    /// Creates a new CRDT vote store with the given node ID.
    ///
    /// # Arguments
    ///
    /// * `store` - The underlying KVStore (wrapped in Arc for sharing)
    /// * `node_id` - Unique identifier for this node (for provenance)
    pub fn new(store: Arc<S>, node_id: [u8; 16]) -> Self {
        Self { inner: GenericVoteStore::new(store.clone()), store, node_id }
    }

    /// Returns a reference to the underlying VoteStore.
    ///
    /// This allows using all standard VoteStore operations.
    pub fn inner(&self) -> &GenericVoteStore<Arc<S>> {
        &self.inner
    }

    /// Gets the vote count state for a specific assertion.
    #[instrument(skip(self))]
    pub async fn get_assertion_state(
        &self,
        assertion_hash: &Hash,
        subject: &str,
    ) -> Result<VoteCountState> {
        let count = self.inner.get_vote_count(assertion_hash, subject).await?;
        let weight = self.inner.get_aggregate_weight(assertion_hash, subject).await?;

        Ok(VoteCountState::new(*assertion_hash, count, weight, self.node_id))
    }

    /// Merges a single assertion's vote state from a remote node.
    ///
    /// Uses G-Counter semantics: takes the maximum of local and remote values.
    #[instrument(skip(self, remote), fields(
        assertion_hash = %hex::encode(remote.assertion_hash),
        remote_count = remote.count,
        remote_weight = remote.weight
    ))]
    pub async fn merge_assertion_state(
        &self,
        subject: &str,
        remote: &VoteCountState,
    ) -> Result<()> {
        let assertion_hex = hex::encode(remote.assertion_hash);

        // Get current local state
        let local_count = self.inner.get_vote_count(&remote.assertion_hash, subject).await?;
        let local_weight = self.inner.get_aggregate_weight(&remote.assertion_hash, subject).await?;

        // Apply G-Counter merge: take max
        if remote.count > local_count {
            let count_key = key_codec::vote_count_key(subject, &assertion_hex);
            // Set to the higher value
            // Note: This is safe because counts only grow in G-Counters
            self.store.put(&count_key, &remote.count.to_le_bytes()).await?;
            debug!(
                old_count = local_count,
                new_count = remote.count,
                "Merged vote count (remote was higher)"
            );
        }

        if remote.weight > local_weight {
            let weight_key = key_codec::vote_weight_key(subject, &assertion_hex);
            // Set to the higher value
            self.store.put(&weight_key, &remote.weight.to_le_bytes()).await?;
            debug!(
                old_weight = local_weight,
                new_weight = remote.weight,
                "Merged aggregate weight (remote was higher)"
            );
        }

        Ok(())
    }
}

/// Aggregate state for all votes under a subject.
///
/// Used for bulk state transfer during initial sync or catch-up.
#[derive(Archive, Deserialize, Serialize, Debug, Clone, PartialEq)]
#[archive(check_bytes)]
pub struct SubjectVoteState {
    /// The subject this state covers.
    pub subject: String,
    /// Vote states for each assertion under this subject.
    pub assertions: Vec<VoteCountState>,
    /// Source node ID.
    pub source_node: [u8; 16],
}

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

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

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

        for (key, count_bytes) in entries {
            // Extract assertion hash from key
            // Key format: {subject}\x00VC:{assertion_hex}
            let key_str = String::from_utf8_lossy(&key);
            if let Some(assertion_hex) = key_str.split(':').next_back() {
                if let Ok(assertion_hash) = hex::decode(assertion_hex) {
                    if assertion_hash.len() == 32 {
                        let hash: Hash = assertion_hash.try_into().map_err(|_| {
                            crate::error::StorageError::Serialization(
                                "Invalid assertion hash".to_string(),
                            )
                        })?;

                        let count = if count_bytes.len() == 8 {
                            u64::from_le_bytes(count_bytes.try_into().map_err(|_| {
                                crate::error::StorageError::Serialization(
                                    "Invalid count bytes".to_string(),
                                )
                            })?)
                        } else {
                            0
                        };

                        // Weight may fail to fetch if store is corrupted; log and use 0.0
                        let weight = match self.inner.get_aggregate_weight(&hash, subject).await {
                            Ok(w) => w,
                            Err(e) => {
                                warn!(
                                    error = %e,
                                    hash = %hex::encode(hash),
                                    "Failed to get aggregate weight, using 0.0"
                                );
                                0.0
                            }
                        };

                        assertions.push(VoteCountState::new(hash, count, weight, self.node_id));
                    }
                }
            }
        }

        Ok(SubjectVoteState { subject: subject.to_string(), assertions, source_node: self.node_id })
    }

    #[instrument(skip(self, remote), fields(subject = %remote.subject, assertion_count = remote.assertions.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(crate::error::StorageError::InputValidation(
                "Subject mismatch in merge".to_string(),
            ));
        }

        for assertion_state in &remote.assertions {
            self.merge_assertion_state(subject, assertion_state).await?;
        }

        debug!(merged_count = remote.assertions.len(), "Merged vote state from remote node");

        Ok(())
    }

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

// Delegate VoteStore trait to inner
#[async_trait]
impl<S: KVStore + 'static> VoteStore for CrdtVoteStore<S> {
    async fn put_vote(&self, vote: &stemedb_core::types::Vote, subject: &str) -> Result<Hash> {
        self.inner.put_vote(vote, subject).await
    }

    async fn get_vote(
        &self,
        assertion_hash: &Hash,
        vote_hash: &Hash,
        subject: &str,
    ) -> Result<Option<stemedb_core::types::Vote>> {
        self.inner.get_vote(assertion_hash, vote_hash, subject).await
    }

    async fn get_votes_for_assertion(
        &self,
        assertion_hash: &Hash,
        subject: &str,
    ) -> Result<Vec<stemedb_core::types::Vote>> {
        self.inner.get_votes_for_assertion(assertion_hash, subject).await
    }

    async fn get_vote_count(&self, assertion_hash: &Hash, subject: &str) -> Result<u64> {
        self.inner.get_vote_count(assertion_hash, subject).await
    }

    async fn get_aggregate_weight(&self, assertion_hash: &Hash, subject: &str) -> Result<f32> {
        self.inner.get_aggregate_weight(assertion_hash, subject).await
    }

    async fn has_votes(&self, assertion_hash: &Hash, subject: &str) -> Result<bool> {
        self.inner.has_votes(assertion_hash, subject).await
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::HybridStore;
    use stemedb_core::types::Vote;
    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_crdt_vote_store_basic() {
        let store = create_test_store().await;
        let node_id = [1u8; 16];
        let crdt = CrdtVoteStore::new(store, node_id);

        let vote = Vote {
            assertion_hash: [1u8; 32],
            agent_id: [2u8; 32],
            weight: 0.8,
            signature: [0u8; 64],
            timestamp: 12345,
            source_url: None,
            observed_context: None,
        };

        // Put a vote
        crdt.put_vote(&vote, "test_subject").await.expect("put_vote");

        // Check count
        let count = crdt.get_vote_count(&[1u8; 32], "test_subject").await.expect("get_count");
        assert_eq!(count, 1);

        // Check weight
        let weight =
            crdt.get_aggregate_weight(&[1u8; 32], "test_subject").await.expect("get_weight");
        assert!((weight - 0.8).abs() < 0.001);
    }

    #[tokio::test]
    async fn test_vote_count_state_merge() {
        let hash = [1u8; 32];
        let node1 = [1u8; 16];
        let node2 = [2u8; 16];

        let state1 = VoteCountState::new(hash, 10, 5.0, node1);
        let state2 = VoteCountState::new(hash, 15, 3.0, node2);

        // Merge: take max of each field
        let merged = state1.merge(&state2);
        assert_eq!(merged.count, 15); // max(10, 15)
        assert_eq!(merged.weight, 5.0); // max(5.0, 3.0)
    }

    #[tokio::test]
    async fn test_crdt_merge_higher_remote() {
        let store1 = create_test_store().await;
        let store2 = create_test_store().await;
        let node1 = [1u8; 16];
        let node2 = [2u8; 16];

        let crdt1 = CrdtVoteStore::new(store1.clone(), node1);
        let crdt2 = CrdtVoteStore::new(store2.clone(), node2);

        // Add votes to node1
        let vote1 = Vote {
            assertion_hash: [1u8; 32],
            agent_id: [2u8; 32],
            weight: 0.5,
            signature: [0u8; 64],
            timestamp: 12345,
            source_url: None,
            observed_context: None,
        };
        crdt1.put_vote(&vote1, "subject").await.expect("put");

        // Add more votes to node2
        for i in 0..5 {
            let vote = Vote {
                assertion_hash: [1u8; 32],
                agent_id: [(i + 10) as u8; 32],
                weight: 0.3,
                signature: [0u8; 64],
                timestamp: 12345 + i as u64,
                source_url: None,
                observed_context: None,
            };
            crdt2.put_vote(&vote, "subject").await.expect("put");
        }

        // Get state from node2
        let state2 = crdt2.get_state("subject").await.expect("get_state");
        assert_eq!(state2.assertions.len(), 1);
        assert_eq!(state2.assertions[0].count, 5);

        // Merge into node1
        crdt1.merge("subject", &state2).await.expect("merge");

        // Node1 should now have higher count
        let count = crdt1.get_vote_count(&[1u8; 32], "subject").await.expect("get");
        assert_eq!(count, 5); // Merged from node2
    }

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

        // Create a state to merge
        let remote_state = SubjectVoteState {
            subject: "test".to_string(),
            assertions: vec![VoteCountState::new([1u8; 32], 10, 5.0, [2u8; 16])],
            source_node: [2u8; 16],
        };

        // Merge once
        crdt.merge("test", &remote_state).await.expect("merge1");
        let count1 = crdt.get_vote_count(&[1u8; 32], "test").await.expect("get");

        // Merge again (should be idempotent)
        crdt.merge("test", &remote_state).await.expect("merge2");
        let count2 = crdt.get_vote_count(&[1u8; 32], "test").await.expect("get");

        assert_eq!(count1, count2);
    }
}
