//! AuditStore implementation backed by a generic KVStore.

use crate::error::Result;
use crate::key_codec;
use crate::traits::KVStore;
use async_trait::async_trait;
use stemedb_core::types::{QueryAudit, QueryId};
use tracing::{debug, instrument};

use super::AuditStore;

/// AuditStore implementation backed by a generic KVStore.
///
/// This implementation maintains an agent index for efficient temporal queries.
pub struct GenericAuditStore<S> {
    store: S,
}

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

    /// Construct the key for an individual audit record.
    pub(crate) fn audit_key(query_id: &QueryId) -> Vec<u8> {
        let query_hex = hex::encode(query_id);
        key_codec::audit_key(&query_hex)
    }

    /// Construct the agent index key.
    pub(crate) fn agent_index_key(
        agent_id: &[u8; 32],
        timestamp: u64,
        query_id: &QueryId,
    ) -> Vec<u8> {
        let agent_hex = hex::encode(agent_id);
        let timestamp_hex = format!("{:016x}", timestamp);
        let query_hex = hex::encode(query_id);
        key_codec::audit_agent_index_key(&agent_hex, &timestamp_hex, &query_hex)
    }

    /// Construct the prefix for scanning all audits for an agent.
    pub(crate) fn agent_full_prefix(agent_id: &[u8; 32]) -> Vec<u8> {
        let agent_hex = hex::encode(agent_id);
        key_codec::audit_agent_prefix(&agent_hex)
    }

    /// Serialize an audit using the canonical serde helpers.
    pub(crate) fn serialize_audit(audit: &QueryAudit) -> Result<Vec<u8>> {
        crate::serde_helpers::serialize(audit)
    }

    /// Deserialize an audit using the canonical serde helpers.
    pub(crate) fn deserialize_audit(data: &[u8]) -> Result<QueryAudit> {
        crate::serde_helpers::deserialize(data)
    }

    /// Extract query_id from an agent index key.
    pub(crate) fn extract_query_id_from_key(key: &[u8]) -> Option<QueryId> {
        // Key format: \x00AUDA:{agent_hex}:{timestamp_hex}:{query_hex}
        let key_str = std::str::from_utf8(key).ok()?;

        // Skip the leading \x00 if present
        let key_content = key_str.strip_prefix('\x00').unwrap_or(key_str);

        let parts: Vec<&str> = key_content.split(':').collect();
        if parts.len() != 4 {
            return None;
        }
        let query_hex = parts[3];
        let bytes = hex::decode(query_hex).ok()?;
        if bytes.len() != 32 {
            return None;
        }
        let mut query_id = [0u8; 32];
        query_id.copy_from_slice(&bytes);
        Some(query_id)
    }

    /// Extract timestamp from an agent index key.
    pub(crate) fn extract_timestamp_from_key(key: &[u8]) -> Option<u64> {
        // Key format: \x00AUDA:{agent_hex}:{timestamp_hex}:{query_hex}
        let key_str = std::str::from_utf8(key).ok()?;

        // Skip the leading \x00 if present
        let key_content = key_str.strip_prefix('\x00').unwrap_or(key_str);

        let parts: Vec<&str> = key_content.split(':').collect();
        if parts.len() != 4 {
            return None;
        }
        u64::from_str_radix(parts[2], 16).ok()
    }
}

#[async_trait]
impl<S: KVStore + 'static> AuditStore for GenericAuditStore<S> {
    #[instrument(skip(self, audit), fields(query_id = %hex::encode(audit.query_id)))]
    async fn put_audit(&self, audit: &QueryAudit) -> Result<QueryId> {
        // Serialize the audit
        let serialized = Self::serialize_audit(audit)?;

        // Store the audit record
        let key = Self::audit_key(&audit.query_id);
        self.store.put(&key, &serialized).await?;

        debug!(
            query_id = %hex::encode(audit.query_id),
            "Stored query audit"
        );

        // Create agent index entry if agent_id is present
        if let Some(agent_id) = &audit.agent_id {
            let index_key = Self::agent_index_key(agent_id, audit.timestamp, &audit.query_id);
            // Store empty value - the key itself contains all the info
            self.store.put(&index_key, &[]).await?;
            debug!(
                agent_id = %hex::encode(agent_id),
                timestamp = audit.timestamp,
                "Created agent audit index"
            );
        }

        Ok(audit.query_id)
    }

    #[instrument(skip(self), fields(query_id = %hex::encode(query_id)))]
    async fn get_audit(&self, query_id: &QueryId) -> Result<Option<QueryAudit>> {
        let key = Self::audit_key(query_id);
        match self.store.get(&key).await? {
            Some(data) => {
                let audit = Self::deserialize_audit(&data)?;
                Ok(Some(audit))
            }
            None => Ok(None),
        }
    }

    #[instrument(skip(self), fields(agent_id = %hex::encode(agent_id), from = from_timestamp, limit))]
    async fn get_audits_for_agent(
        &self,
        agent_id: &[u8; 32],
        from_timestamp: u64,
        to_timestamp: Option<u64>,
        limit: usize,
    ) -> Result<Vec<QueryAudit>> {
        // Scan the agent index
        let prefix = Self::agent_full_prefix(agent_id);
        let entries = self.store.scan_prefix(&prefix).await?;

        let mut audits = Vec::new();

        for (key, _value) in entries {
            // Extract timestamp and filter
            let Some(timestamp) = Self::extract_timestamp_from_key(&key) else {
                continue;
            };

            // Check time range
            if timestamp < from_timestamp {
                continue;
            }
            if let Some(to) = to_timestamp {
                if timestamp > to {
                    continue;
                }
            }

            // Extract query_id and fetch the full audit
            let Some(query_id) = Self::extract_query_id_from_key(&key) else {
                continue;
            };

            if let Some(audit) = self.get_audit(&query_id).await? {
                audits.push(audit);
                // Early termination once we have enough (before sorting)
                // Note: We collect more than limit to allow proper sorting
                // A future optimization could use timestamp-ordered keys
                if audits.len() >= limit * 2 {
                    break;
                }
            }
        }

        // Sort by timestamp ascending
        audits.sort_by_key(|a| a.timestamp);

        // Apply limit after sorting
        audits.truncate(limit);

        debug!(count = audits.len(), "Retrieved audits for agent");
        Ok(audits)
    }

    #[instrument(skip(self), fields(limit))]
    async fn list_recent_audits(&self, limit: usize) -> Result<Vec<QueryAudit>> {
        let prefix = key_codec::audit_scan_prefix();
        let entries = self.store.scan_prefix(&prefix).await?;

        let mut audits = Vec::with_capacity(entries.len().min(limit));

        for (_key, data) in entries {
            match Self::deserialize_audit(&data) {
                Ok(audit) => audits.push(audit),
                Err(e) => {
                    debug!(error = %e, "Skipping malformed audit record");
                }
            }
        }

        // Sort by timestamp descending (most recent first)
        audits.sort_by(|a, b| b.timestamp.cmp(&a.timestamp));

        // Apply limit
        audits.truncate(limit);

        debug!(count = audits.len(), "Listed recent audits");
        Ok(audits)
    }

    #[instrument(skip(self), fields(agent_id = %hex::encode(agent_id)))]
    async fn has_audits_for_agent(&self, agent_id: &[u8; 32]) -> Result<bool> {
        let prefix = Self::agent_full_prefix(agent_id);
        let entries = self.store.scan_prefix(&prefix).await?;
        Ok(!entries.is_empty())
    }
}
