//! GenericPatternAggregateStore implementation backed by a generic KVStore.

use async_trait::async_trait;
use tracing::{debug, instrument};

use super::{PatternAggregate, PatternAggregateStore};
use crate::error::{Result, StorageError};
use crate::key_codec;
use crate::traits::KVStore;

/// PatternAggregateStore implementation backed by a generic KVStore.
pub struct GenericPatternAggregateStore<S> {
    store: S,
}

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

    /// Serialize an aggregate using serde_json.
    fn serialize_aggregate(agg: &PatternAggregate) -> Result<Vec<u8>> {
        serde_json::to_vec(agg).map_err(|e| StorageError::Serialization(e.to_string()))
    }

    /// Deserialize an aggregate using serde_json.
    fn deserialize_aggregate(data: &[u8]) -> Result<PatternAggregate> {
        serde_json::from_slice(data).map_err(|e| StorageError::Serialization(e.to_string()))
    }
}

#[async_trait]
impl<S: KVStore + 'static> PatternAggregateStore for GenericPatternAggregateStore<S> {
    #[instrument(skip(self), fields(subject, predicate, value_hash, project_hash))]
    async fn record_observation(
        &self,
        subject: &str,
        predicate: &str,
        value_hash: &str,
        value_display: &str,
        project_hash: &str,
        timestamp: u64,
    ) -> Result<PatternAggregate> {
        // 1. Check if this project has already reported this pattern
        let project_set_key =
            key_codec::pattern_project_set_key(subject, predicate, value_hash, project_hash);
        let is_new_project = self.store.get(&project_set_key).await?.is_none();

        // 2. If new project, record it in the project set
        if is_new_project {
            self.store.put(&project_set_key, &[1u8]).await?;
        }

        // 3. Atomically increment observation count (always)
        let obs_count_key =
            key_codec::pattern_observation_count_key(subject, predicate, value_hash);
        let obs_count = self.store.fetch_and_add_u64(&obs_count_key, 1).await?;

        // 4. Atomically increment project count (only if new project)
        let proj_count_key = key_codec::pattern_project_count_key(subject, predicate, value_hash);
        let proj_count = if is_new_project {
            self.store.fetch_and_add_u64(&proj_count_key, 1).await?
        } else {
            // Read current count
            match self.store.get(&proj_count_key).await? {
                Some(bytes) if bytes.len() == 8 => {
                    let arr: [u8; 8] = bytes.try_into().map_err(|_| {
                        StorageError::Serialization("Invalid project count bytes".to_string())
                    })?;
                    u64::from_le_bytes(arr)
                }
                _ => 0,
            }
        };

        // 5. Update or create the aggregate record
        let agg_key = key_codec::pattern_aggregate_key(subject, predicate, value_hash);
        let agg = match self.store.get(&agg_key).await? {
            Some(data) => {
                let mut existing = Self::deserialize_aggregate(&data)?;
                existing.observation_count = obs_count;
                existing.project_count = proj_count;
                existing.last_seen = timestamp;
                existing
            }
            None => PatternAggregate {
                subject: subject.to_string(),
                predicate: predicate.to_string(),
                value_hash: value_hash.to_string(),
                value_display: value_display.to_string(),
                project_count: proj_count,
                observation_count: obs_count,
                first_seen: timestamp,
                last_seen: timestamp,
            },
        };

        // Store the updated aggregate
        let serialized = Self::serialize_aggregate(&agg)?;
        self.store.put(&agg_key, &serialized).await?;

        debug!(
            project_count = proj_count,
            observation_count = obs_count,
            is_new_project,
            "Recorded observation"
        );

        Ok(agg)
    }

    #[instrument(skip(self), fields(subject, predicate, value_hash))]
    async fn get_aggregate(
        &self,
        subject: &str,
        predicate: &str,
        value_hash: &str,
    ) -> Result<Option<PatternAggregate>> {
        let key = key_codec::pattern_aggregate_key(subject, predicate, value_hash);
        match self.store.get(&key).await? {
            Some(data) => {
                let agg = Self::deserialize_aggregate(&data)?;
                Ok(Some(agg))
            }
            None => Ok(None),
        }
    }

    #[instrument(skip(self), fields(min_projects, limit))]
    async fn get_popular_patterns(
        &self,
        min_projects: u64,
        limit: usize,
    ) -> Result<Vec<PatternAggregate>> {
        // Scan all pattern aggregates using a global prefix scan
        // Note: In production, we'd want an index sorted by project_count
        // For now, we scan all and filter

        // We need to scan with a prefix that catches all subjects
        // Since subjects start with "code://" etc., we'll scan each known scheme
        let schemes = ["code://", "vendor://", "rfc://", "owasp://"];

        let mut all_patterns = Vec::new();

        for scheme in schemes {
            // Scan all subjects with this scheme
            let prefix = scheme.as_bytes().to_vec();
            let entries = self.store.scan_prefix(&prefix).await?;

            for (key, data) in entries {
                // Check if this is a pattern aggregate key (contains PA:)
                if let Ok(key_str) = std::str::from_utf8(&key) {
                    if key_str.contains("\x00PA:") {
                        if let Ok(agg) = Self::deserialize_aggregate(&data) {
                            if agg.project_count >= min_projects {
                                all_patterns.push(agg);
                            }
                        }
                    }
                }
            }
        }

        // Sort by project_count descending
        all_patterns.sort_by(|a, b| b.project_count.cmp(&a.project_count));

        // Limit results
        all_patterns.truncate(limit);

        debug!(found = all_patterns.len(), "Retrieved popular patterns");

        Ok(all_patterns)
    }

    #[instrument(skip(self), fields(subject_prefix, min_projects, limit))]
    async fn get_patterns_for_subject_prefix(
        &self,
        subject_prefix: &str,
        min_projects: u64,
        limit: usize,
    ) -> Result<Vec<PatternAggregate>> {
        let prefix = subject_prefix.as_bytes().to_vec();
        let entries = self.store.scan_prefix(&prefix).await?;

        let mut patterns = Vec::new();

        for (key, data) in entries {
            // Check if this is a pattern aggregate key (contains PA:)
            if let Ok(key_str) = std::str::from_utf8(&key) {
                if key_str.contains("\x00PA:") {
                    if let Ok(agg) = Self::deserialize_aggregate(&data) {
                        if agg.project_count >= min_projects {
                            patterns.push(agg);
                        }
                    }
                }
            }
        }

        // Sort by project_count descending
        patterns.sort_by(|a, b| b.project_count.cmp(&a.project_count));

        // Limit results
        patterns.truncate(limit);

        debug!(found = patterns.len(), "Retrieved patterns for prefix");

        Ok(patterns)
    }
}
