//! ApiKeyStore implementation backed by a generic KVStore.

use super::{ApiKeyRecord, ApiKeyStore, RateLimitResult, DEFAULT_API_KEY_RATE_LIMIT};
use crate::error::{Result, StorageError};
use crate::key_codec;
use crate::traits::KVStore;
use async_trait::async_trait;
use tracing::{debug, instrument, warn};

/// ApiKeyStore implementation backed by a generic KVStore.
#[derive(Clone)]
pub struct GenericApiKeyStore<S> {
    store: S,
    default_rate_limit: u64,
}

impl<S: KVStore> GenericApiKeyStore<S> {
    /// Create a new ApiKeyStore with default configuration.
    pub fn new(store: S) -> Self {
        Self { store, default_rate_limit: DEFAULT_API_KEY_RATE_LIMIT }
    }

    /// Create a new ApiKeyStore with custom default rate limit.
    pub fn with_rate_limit(store: S, default_rate_limit: u64) -> Self {
        Self { store, default_rate_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 an ApiKeyRecord.
    pub(crate) fn serialize_record(record: &ApiKeyRecord) -> Result<Vec<u8>> {
        crate::serde_helpers::serialize(record)
    }

    /// Deserialize an ApiKeyRecord.
    pub(crate) fn deserialize_record(data: &[u8]) -> Result<ApiKeyRecord> {
        crate::serde_helpers::deserialize(data)
    }

    /// Get the default rate limit.
    pub fn default_rate_limit(&self) -> u64 {
        self.default_rate_limit
    }
}

#[async_trait]
impl<S: KVStore + 'static> ApiKeyStore for GenericApiKeyStore<S> {
    #[instrument(skip(self, record), fields(key_prefix = %record.key_prefix, role = %record.role))]
    async fn put_key(&self, key_hash: &[u8; 32], record: &ApiKeyRecord) -> Result<()> {
        let key = key_codec::api_key_record_key(&hex::encode(key_hash));
        let serialized = Self::serialize_record(record)?;
        self.store.put(&key, &serialized).await?;
        debug!("Stored API key");
        Ok(())
    }

    #[instrument(skip(self), fields(key_hash_prefix = %hex::encode(&key_hash[..4])))]
    async fn get_key_by_hash(&self, key_hash: &[u8; 32]) -> Result<Option<ApiKeyRecord>> {
        let key = key_codec::api_key_record_key(&hex::encode(key_hash));
        match self.store.get(&key).await? {
            Some(data) => {
                let record = Self::deserialize_record(&data)?;
                Ok(Some(record))
            }
            None => Ok(None),
        }
    }

    #[instrument(skip(self), fields(key_hash_prefix = %hex::encode(&key_hash[..4])))]
    async fn validate_key(&self, key_hash: &[u8; 32], now: u64) -> Result<Option<ApiKeyRecord>> {
        let record = match self.get_key_by_hash(key_hash).await? {
            Some(r) => r,
            None => {
                debug!("API key not found");
                return Ok(None);
            }
        };

        if !record.enabled {
            debug!(label = %record.label, "API key is disabled");
            return Ok(None);
        }

        if record.is_expired(now) {
            debug!(label = %record.label, expires_at = ?record.expires_at, "API key has expired");
            return Ok(None);
        }

        debug!(label = %record.label, role = %record.role, "API key validated");
        Ok(Some(record))
    }

    #[instrument(skip(self))]
    async fn list_keys(&self) -> Result<Vec<ApiKeyRecord>> {
        let prefix = key_codec::api_key_scan_prefix();
        let pairs = self.store.scan_prefix(&prefix).await?;

        let mut records = Vec::with_capacity(pairs.len());
        for (_key, value) in pairs {
            match Self::deserialize_record(&value) {
                Ok(record) => records.push(record),
                Err(e) => {
                    warn!(error = %e, "Failed to deserialize API key record, skipping");
                }
            }
        }

        debug!(count = records.len(), "Listed API keys");
        Ok(records)
    }

    #[instrument(skip(self), fields(key_hash_prefix = %hex::encode(&key_hash[..4])))]
    async fn delete_key(&self, key_hash: &[u8; 32]) -> Result<()> {
        let key = key_codec::api_key_record_key(&hex::encode(key_hash));
        self.store.delete(&key).await?;
        debug!("Deleted API key");
        Ok(())
    }

    #[instrument(skip(self), fields(key_hash_prefix = %hex::encode(&key_hash[..4])))]
    async fn touch_key(&self, key_hash: &[u8; 32], timestamp: u64) -> Result<()> {
        let key = key_codec::api_key_record_key(&hex::encode(key_hash));
        let data = self.store.get(&key).await?.ok_or_else(|| {
            StorageError::NotFound(format!("API key not found: {}", hex::encode(&key_hash[..8])))
        })?;

        let mut record = Self::deserialize_record(&data)?;
        record.last_used_at = timestamp;

        let serialized = Self::serialize_record(&record)?;
        self.store.put(&key, &serialized).await?;
        debug!("Touched API key");
        Ok(())
    }

    #[instrument(skip(self), fields(key_hash_prefix = %hex::encode(&key_hash[..4]), limit))]
    async fn check_rate_limit(
        &self,
        key_hash: &[u8; 32],
        limit: u64,
        timestamp: u64,
    ) -> Result<RateLimitResult> {
        let window_start = Self::hour_window(timestamp);
        let reset_at = Self::reset_timestamp(timestamp);

        let key = key_codec::api_key_rate_key(&hex::encode(key_hash), window_start);

        // Get current usage for this window
        let current_usage = match self.store.get(&key).await? {
            Some(bytes) if bytes.len() == 8 => {
                let arr: [u8; 8] = bytes
                    .try_into()
                    .map_err(|_| StorageError::Serialization("Invalid rate bytes".to_string()))?;
                u64::from_le_bytes(arr)
            }
            _ => 0,
        };

        // Check if would exceed limit
        if current_usage >= limit {
            warn!(
                current_usage,
                limit,
                key_hash_prefix = %hex::encode(&key_hash[..4]),
                "API key rate limited"
            );
            return Ok(RateLimitResult { allowed: false, remaining: 0, limit, reset_at });
        }

        // Record this request
        let new_usage = current_usage.saturating_add(1);
        self.store.put(&key, &new_usage.to_le_bytes()).await?;

        let remaining = limit.saturating_sub(new_usage);
        debug!(current_usage = new_usage, remaining, "Rate limit check passed");

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

    #[instrument(skip(self), fields(key_hash_prefix = %hex::encode(&key_hash[..4]), enabled))]
    async fn set_enabled(&self, key_hash: &[u8; 32], enabled: bool) -> Result<()> {
        let key = key_codec::api_key_record_key(&hex::encode(key_hash));
        let data = self.store.get(&key).await?.ok_or_else(|| {
            StorageError::NotFound(format!("API key not found: {}", hex::encode(&key_hash[..8])))
        })?;

        let mut record = Self::deserialize_record(&data)?;
        record.enabled = enabled;

        let serialized = Self::serialize_record(&record)?;
        self.store.put(&key, &serialized).await?;
        debug!(label = %record.label, "Updated API key enabled state");
        Ok(())
    }
}
