//! QuotaStore implementation backed by a generic KVStore.

use super::{
    CostConfig, OperationType, QuotaCheckResult, QuotaRecord, QuotaStore, DEFAULT_QUOTA_LIMIT,
};
use crate::error::{Result, StorageError};
use crate::key_codec;
use crate::traits::KVStore;
use async_trait::async_trait;
use tracing::{debug, instrument, warn};

/// QuotaStore implementation backed by a generic KVStore.
#[derive(Clone)]
pub struct GenericQuotaStore<S> {
    store: S,
    cost_config: CostConfig,
    default_limit: u64,
}

impl<S: KVStore> GenericQuotaStore<S> {
    /// Create a new QuotaStore with default configuration.
    pub fn new(store: S) -> Self {
        Self { store, cost_config: CostConfig::default(), default_limit: DEFAULT_QUOTA_LIMIT }
    }

    /// Create a new QuotaStore with custom configuration.
    pub fn with_config(store: S, cost_config: CostConfig, default_limit: u64) -> Self {
        Self { store, cost_config, default_limit }
    }

    /// Get the hour window start for a given timestamp.
    ///
    /// Truncates to the start of the hour (seconds since epoch, divisible by 3600).
    pub(crate) fn hour_window(timestamp: u64) -> u64 {
        (timestamp / 3600) * 3600
    }

    /// Get the reset timestamp (start of next hour).
    pub(crate) fn reset_timestamp(timestamp: u64) -> u64 {
        Self::hour_window(timestamp) + 3600
    }

    /// Serialize a QuotaRecord using the canonical serde helpers.
    pub(crate) fn serialize_record(record: &QuotaRecord) -> Result<Vec<u8>> {
        crate::serde_helpers::serialize(record)
    }

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

    /// Calculate cost for an operation.
    fn calculate_cost(&self, operation: OperationType, payload_bytes: usize) -> u64 {
        match operation {
            OperationType::Assert => self.cost_config.cost_for_assert(payload_bytes),
            OperationType::Vote => self.cost_config.cost_for_vote(payload_bytes),
            OperationType::Query { lens_count } => {
                self.cost_config.cost_for_query(lens_count, payload_bytes)
            }
        }
    }
}

#[async_trait]
impl<S: KVStore + 'static> QuotaStore for GenericQuotaStore<S> {
    #[instrument(skip(self), fields(agent_id = %hex::encode(agent_id), operation = ?operation, payload_bytes))]
    async fn check_and_record(
        &self,
        agent_id: &[u8; 32],
        operation: OperationType,
        payload_bytes: usize,
        timestamp: u64,
    ) -> Result<QuotaCheckResult> {
        let window_start = Self::hour_window(timestamp);
        let reset_at = Self::reset_timestamp(timestamp);
        let cost = self.calculate_cost(operation, payload_bytes);
        let limit = self.get_quota_limit(agent_id).await?;

        // Get or create quota record for this window
        let key = key_codec::quota_key(&hex::encode(agent_id), window_start);
        let mut record = match self.store.get(&key).await? {
            Some(data) => Self::deserialize_record(&data)?,
            None => QuotaRecord::new(*agent_id, window_start),
        };

        // Check if operation would exceed quota
        if record.would_exceed(cost, limit) {
            warn!(
                agent_id = %hex::encode(agent_id),
                cost,
                tokens_used = record.tokens_used,
                limit,
                "Quota exceeded"
            );
            return Ok(QuotaCheckResult {
                allowed: false,
                remaining: record.remaining(limit),
                limit,
                reset_at,
                cost,
            });
        }

        // Record the cost
        record.record(cost, timestamp);
        let remaining = record.remaining(limit);

        // Persist
        let serialized = Self::serialize_record(&record)?;
        self.store.put(&key, &serialized).await?;

        debug!(cost, tokens_used = record.tokens_used, remaining, "Quota recorded");

        Ok(QuotaCheckResult { allowed: true, remaining, limit, reset_at, cost })
    }

    #[instrument(skip(self), fields(agent_id = %hex::encode(agent_id)))]
    async fn get_quota_status(
        &self,
        agent_id: &[u8; 32],
        timestamp: u64,
    ) -> Result<QuotaCheckResult> {
        let window_start = Self::hour_window(timestamp);
        let reset_at = Self::reset_timestamp(timestamp);
        let limit = self.get_quota_limit(agent_id).await?;

        let key = key_codec::quota_key(&hex::encode(agent_id), window_start);
        let record = match self.store.get(&key).await? {
            Some(data) => Self::deserialize_record(&data)?,
            None => QuotaRecord::new(*agent_id, window_start),
        };

        Ok(QuotaCheckResult {
            allowed: true, // Status check doesn't consume quota
            remaining: record.remaining(limit),
            limit,
            reset_at,
            cost: 0,
        })
    }

    #[instrument(skip(self), fields(agent_id = %hex::encode(agent_id), limit))]
    async fn set_quota_limit(&self, agent_id: &[u8; 32], limit: u64) -> Result<()> {
        let key = key_codec::quota_limit_key(&hex::encode(agent_id));
        self.store.put(&key, &limit.to_le_bytes()).await?;
        debug!("Set custom quota limit");
        Ok(())
    }

    #[instrument(skip(self), fields(agent_id = %hex::encode(agent_id)))]
    async fn get_quota_limit(&self, agent_id: &[u8; 32]) -> Result<u64> {
        let key = key_codec::quota_limit_key(&hex::encode(agent_id));
        match self.store.get(&key).await? {
            Some(bytes) if bytes.len() == 8 => {
                let arr: [u8; 8] = bytes
                    .try_into()
                    .map_err(|_| StorageError::Serialization("Invalid limit bytes".to_string()))?;
                Ok(u64::from_le_bytes(arr))
            }
            _ => Ok(self.default_limit),
        }
    }

    fn cost_config(&self) -> &CostConfig {
        &self.cost_config
    }
}
