//! GenericVoteStore implementation backed by a generic KVStore.

use crate::error::{Result, StorageError};
use crate::key_codec;
use crate::traits::KVStore;
use async_trait::async_trait;
use stemedb_core::types::{Hash, Vote};
use tracing::{debug, instrument};

use super::VoteStore;

/// VoteStore implementation backed by a generic KVStore.
///
/// This implementation maintains caches for vote counts and aggregate weights
/// to enable O(1) lookups for the Materializer.
pub struct GenericVoteStore<S> {
    store: S,
}

impl<S: KVStore> GenericVoteStore<S> {
    /// Create a new VoteStore backed by the given KVStore.
    pub fn new(store: S) -> Self {
        Self { store }
    }

    /// Serialize a vote using the canonical serde helpers.
    fn serialize_vote(vote: &Vote) -> Result<Vec<u8>> {
        crate::serde_helpers::serialize(vote)
    }

    /// Deserialize a vote using the canonical serde helpers.
    fn deserialize_vote(data: &[u8]) -> Result<Vote> {
        crate::serde_helpers::deserialize(data)
    }
}

#[async_trait]
impl<S: KVStore + 'static> VoteStore for GenericVoteStore<S> {
    #[instrument(skip(self, vote), fields(assertion_hash = %hex::encode(vote.assertion_hash), weight = vote.weight, subject))]
    async fn put_vote(&self, vote: &Vote, subject: &str) -> Result<Hash> {
        // Serialize the vote
        let serialized = Self::serialize_vote(vote)?;

        // Compute content-addressed hash
        let vote_hash_bytes = blake3::hash(&serialized);
        let vote_hash: Hash = *vote_hash_bytes.as_bytes();

        // Store the vote using subject-prefixed key
        let assertion_hex = hex::encode(vote.assertion_hash);
        let vote_hex = hex::encode(vote_hash);
        let key = key_codec::vote_key(subject, &assertion_hex, &vote_hex);
        self.store.put(&key, &serialized).await?;

        debug!(
            vote_hash = %vote_hex,
            "Stored vote"
        );

        // Update vote count cache (atomic increment - prevents race conditions)
        let count_key = key_codec::vote_count_key(subject, &assertion_hex);
        let new_count = self.store.fetch_and_add_u64(&count_key, 1).await?;

        // Update aggregate weight cache (atomic CAS - prevents race conditions)
        let weight_key = key_codec::vote_weight_key(subject, &assertion_hex);
        let vote_weight = vote.weight;
        let new_weight = self
            .store
            .compare_and_swap_f32(&weight_key, move |current| current + vote_weight)
            .await?;

        debug!(new_count, new_weight, "Updated vote caches (atomic)");

        Ok(vote_hash)
    }

    #[instrument(skip(self), fields(assertion_hash = %hex::encode(assertion_hash), vote_hash = %hex::encode(vote_hash), subject))]
    async fn get_vote(
        &self,
        assertion_hash: &Hash,
        vote_hash: &Hash,
        subject: &str,
    ) -> Result<Option<Vote>> {
        let assertion_hex = hex::encode(assertion_hash);
        let vote_hex = hex::encode(vote_hash);
        let key = key_codec::vote_key(subject, &assertion_hex, &vote_hex);
        match self.store.get(&key).await? {
            Some(data) => {
                let vote = Self::deserialize_vote(&data)?;
                Ok(Some(vote))
            }
            None => Ok(None),
        }
    }

    #[instrument(skip(self), fields(assertion_hash = %hex::encode(assertion_hash), subject))]
    async fn get_votes_for_assertion(
        &self,
        assertion_hash: &Hash,
        subject: &str,
    ) -> Result<Vec<Vote>> {
        let assertion_hex = hex::encode(assertion_hash);
        let prefix = key_codec::vote_scan_prefix(subject, &assertion_hex);
        let entries = self.store.scan_prefix(&prefix).await?;

        let mut votes = Vec::with_capacity(entries.len());
        for (_key, data) in entries {
            let vote = Self::deserialize_vote(&data)?;
            votes.push(vote);
        }

        debug!(count = votes.len(), "Retrieved votes for assertion");
        Ok(votes)
    }

    #[instrument(skip(self), fields(assertion_hash = %hex::encode(assertion_hash), subject))]
    async fn get_vote_count(&self, assertion_hash: &Hash, subject: &str) -> Result<u64> {
        let assertion_hex = hex::encode(assertion_hash);
        let key = key_codec::vote_count_key(subject, &assertion_hex);
        match self.store.get(&key).await? {
            Some(bytes) if bytes.len() == 8 => {
                let arr: [u8; 8] = bytes.try_into().map_err(|_| {
                    StorageError::Serialization("Invalid vote count bytes".to_string())
                })?;
                Ok(u64::from_le_bytes(arr))
            }
            _ => Ok(0),
        }
    }

    #[instrument(skip(self), fields(assertion_hash = %hex::encode(assertion_hash), subject))]
    async fn get_aggregate_weight(&self, assertion_hash: &Hash, subject: &str) -> Result<f32> {
        let assertion_hex = hex::encode(assertion_hash);
        let key = key_codec::vote_weight_key(subject, &assertion_hex);
        match self.store.get(&key).await? {
            Some(bytes) if bytes.len() == 4 => {
                let arr: [u8; 4] = bytes
                    .try_into()
                    .map_err(|_| StorageError::Serialization("Invalid weight bytes".to_string()))?;
                Ok(f32::from_le_bytes(arr))
            }
            _ => Ok(0.0),
        }
    }

    #[instrument(skip(self), fields(assertion_hash = %hex::encode(assertion_hash), subject))]
    async fn has_votes(&self, assertion_hash: &Hash, subject: &str) -> Result<bool> {
        let count = self.get_vote_count(assertion_hash, subject).await?;
        Ok(count > 0)
    }
}
