alien-bindings 3.3.5

use crate::error::{ErrorData, Result};
use crate::traits::{Binding, Kv, PutOptions, ScanResult};
use alien_azure_clients::tables::{
    AzureTableStorageClient, EntityQueryContinuation, EntityQueryOptions, TableEntity,
    TableStorageApi,
};
use alien_error::{AlienError, Context, IntoAlienError};
use async_trait::async_trait;
use base64::{engine::general_purpose::STANDARD as BASE64, Engine as _};
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::collections::HashMap;
use std::fmt::{Debug, Formatter};

use super::{validate_key, validate_value};

/// Convert a KV operation to a Table Storage entity
/// This only base64 encodes the raw bytes when creating the properties map, not in memory
fn create_table_entity(
    partition_key: String,
    row_key: String,
    value: &[u8],
    expires_at: Option<DateTime<Utc>>,
) -> TableEntity {
    let mut properties = HashMap::new();

    // Base64 encode the raw bytes only when storing in the properties map
    // This keeps the original 32KB limit valid since we're not storing the encoded version in memory
    properties.insert("Value".to_string(), Value::String(BASE64.encode(value)));

    // Store creation timestamp
    properties.insert(
        "CreatedAt".to_string(),
        Value::String(Utc::now().to_rfc3339()),
    );

    // Store expiration timestamp if provided
    if let Some(expiry) = expires_at {
        properties.insert("ExpiresAt".to_string(), Value::String(expiry.to_rfc3339()));
    }

    TableEntity {
        partition_key,
        row_key,
        timestamp: None, // Azure will set this
        properties,
    }
}

/// Extract KV value from Table Storage entity
fn extract_value_from_entity(entity: &TableEntity) -> Result<Vec<u8>> {
    let value_str = entity
        .properties
        .get("Value")
        .and_then(|v| v.as_str())
        .ok_or_else(|| {
            AlienError::new(ErrorData::InvalidInput {
                operation_context: "Azure Table Storage KV extract value".to_string(),
                details: "Entity missing Value property or not a string".to_string(),
                field_name: Some("Value".to_string()),
            })
        })?;

    // Decode base64 value
    BASE64
        .decode(value_str)
        .into_alien_error()
        .context(ErrorData::InvalidInput {
            operation_context: "Azure Table Storage KV extract value".to_string(),
            details: "Failed to decode base64 value".to_string(),
            field_name: Some("Value".to_string()),
        })
}

/// Check if entity has expired based on TTL
fn is_entity_expired(entity: &TableEntity) -> bool {
    if let Some(expires_at_value) = entity.properties.get("ExpiresAt") {
        if let Some(expires_at_str) = expires_at_value.as_str() {
            if let Ok(expires_at) = DateTime::parse_from_rfc3339(expires_at_str) {
                return Utc::now() > expires_at.with_timezone(&Utc);
            }
        }
    }
    false
}

/// Cursor state for pagination across partitions
#[derive(Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
struct CursorState {
    version: u8,
    prefix: String,
    current_partition: u32,
    continuation: Option<EntityQueryContinuation>,
}

/// Azure Table Storage implementation of the KV trait
pub struct AzureTableStorageKv {
    client: AzureTableStorageClient,
    resource_group_name: String,
    account_name: String,
    table_name: String,
    num_partitions: u32,
}

impl Debug for AzureTableStorageKv {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("AzureTableStorageKv")
            .field("resource_group_name", &self.resource_group_name)
            .field("account_name", &self.account_name)
            .field("table_name", &self.table_name)
            .field("num_partitions", &self.num_partitions)
            .finish()
    }
}

impl AzureTableStorageKv {
    pub fn new(
        client: AzureTableStorageClient,
        resource_group_name: String,
        account_name: String,
        table_name: String,
    ) -> Self {
        Self {
            client,
            resource_group_name,
            account_name,
            table_name,
            num_partitions: 16, // 16 partitions for load distribution
        }
    }

    /// Creates a hash bucket for load distribution
    fn hash_bucket(&self, key: &str) -> u32 {
        use std::collections::hash_map::DefaultHasher;
        use std::hash::{Hash, Hasher};

        let mut hasher = DefaultHasher::new();
        key.hash(&mut hasher);
        hasher.finish() as u32 % self.num_partitions
    }

    /// Splits key into partition key and row key
    /// Attempt to take over a logically-expired row after a conditional
    /// insert conflict. Returns `Ok(true)` only when THIS caller replaced
    /// the expired row (If-Match on the read ETag); a live row, a lost
    /// race, or a row deleted in between all resolve to `Ok(false)`.
    async fn try_take_over_expired(
        &self,
        key: &str,
        entity: &alien_azure_clients::azure::tables::TableEntity,
    ) -> Result<bool> {
        use alien_client_core::ErrorData as CloudErrorData;

        let (partition_key, row_key) = self.split_key(key);
        let existing = match self
            .client
            .get_entity(
                &self.resource_group_name,
                &self.account_name,
                &self.table_name,
                &partition_key,
                &row_key,
                None,
            )
            .await
        {
            Ok(existing) => existing,
            // Deleted between the conflict and this read: treat as lost —
            // the caller's next attempt sees a clean slate.
            Err(e)
                if matches!(
                    e.error.as_ref(),
                    Some(CloudErrorData::RemoteResourceNotFound { .. })
                ) =>
            {
                return Ok(false);
            }
            Err(e) => {
                return Err(crate::error::map_cloud_client_error(
                    e,
                    format!("Failed to read existing entity for key '{}'", key),
                    Some(key.to_string()),
                ));
            }
        };

        if !is_entity_expired(&existing) {
            return Ok(false);
        }
        let Some(etag) = existing
            .properties
            .get("odata.etag")
            .and_then(|value| value.as_str())
        else {
            // No ETag on the read entity — cannot replace safely.
            return Ok(false);
        };

        match self
            .client
            .update_entity(
                &self.resource_group_name,
                &self.account_name,
                &self.table_name,
                &partition_key,
                &row_key,
                entity,
                Some(alien_azure_clients::azure::tables::ETag::from(etag)),
            )
            .await
        {
            Ok(_) => Ok(true),
            // 409/412: someone else replaced or deleted it first.
            Err(e)
                if matches!(
                    e.error.as_ref(),
                    Some(CloudErrorData::RemoteResourceConflict { .. })
                        | Some(CloudErrorData::RemoteResourceNotFound { .. })
                ) =>
            {
                Ok(false)
            }
            Err(e) => Err(crate::error::map_cloud_client_error(
                e,
                format!("Failed to take over expired entity for key '{}'", key),
                Some(key.to_string()),
            )),
        }
    }

    fn split_key(&self, key: &str) -> (String, String) {
        // Use hash-based partitioning for load distribution
        let partition_key = format!("p{}", self.hash_bucket(key));
        (partition_key, key.to_string())
    }

    /// Combines partition key and row key back to original key
    fn combine_key(&self, _partition_key: &str, row_key: &str) -> String {
        row_key.to_string() // Row key contains the original key
    }

    /// Encodes cursor state as base64url JSON for safe HTTP transmission
    fn encode_cursor(&self, state: &CursorState) -> Result<String> {
        let json =
            serde_json::to_vec(state)
                .into_alien_error()
                .context(ErrorData::InvalidInput {
                    operation_context: "Azure Table Storage KV cursor encoding".to_string(),
                    details: "Failed to serialize cursor state".to_string(),
                    field_name: Some("cursor".to_string()),
                })?;
        Ok(base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(json))
    }

    /// Decodes cursor state from base64url JSON
    fn decode_cursor(&self, prefix: &str, cursor: &str) -> Result<CursorState> {
        let decoded = base64::engine::general_purpose::URL_SAFE_NO_PAD
            .decode(cursor)
            .into_alien_error()
            .context(ErrorData::InvalidInput {
                operation_context: "Azure Table Storage KV cursor decoding".to_string(),
                details: "Invalid cursor encoding".to_string(),
                field_name: Some("cursor".to_string()),
            })?;
        let json =
            String::from_utf8(decoded)
                .into_alien_error()
                .context(ErrorData::InvalidInput {
                    operation_context: "Azure Table Storage KV cursor decoding".to_string(),
                    details: "Invalid cursor UTF-8".to_string(),
                    field_name: Some("cursor".to_string()),
                })?;
        let state: CursorState =
            serde_json::from_str(&json)
                .into_alien_error()
                .context(ErrorData::InvalidInput {
                    operation_context: "Azure Table Storage KV cursor decoding".to_string(),
                    details: "Invalid cursor JSON".to_string(),
                    field_name: Some("cursor".to_string()),
                })?;
        if state.version != 1
            || state.prefix != prefix
            || state.current_partition >= self.num_partitions
        {
            return Err(AlienError::new(ErrorData::InvalidInput {
                operation_context: "Azure Table Storage KV cursor validation".to_string(),
                details: "Cursor does not belong to this prefix scan".to_string(),
                field_name: Some("cursor".to_string()),
            }));
        }
        Ok(state)
    }
}

impl Binding for AzureTableStorageKv {}

#[async_trait]
impl Kv for AzureTableStorageKv {
    async fn get(&self, key: &str) -> Result<Option<Vec<u8>>> {
        validate_key(key)?;

        let (partition_key, row_key) = self.split_key(key);

        match self
            .client
            .get_entity(
                &self.resource_group_name,
                &self.account_name,
                &self.table_name,
                &partition_key,
                &row_key,
                None,
            )
            .await
        {
            Ok(entity) => {
                // Check if TTL has expired (client-side filtering)
                if is_entity_expired(&entity) {
                    return Ok(None); // Expired
                }

                let value = extract_value_from_entity(&entity)?;
                Ok(Some(value))
            }
            Err(e) => {
                use alien_client_core::ErrorData as CloudErrorData;
                match e.error.as_ref() {
                    Some(CloudErrorData::RemoteResourceNotFound { .. }) => Ok(None),
                    _ => Err(crate::error::map_cloud_client_error(
                        e,
                        format!("Failed to get entity for key '{}'", key),
                        Some(key.to_string()),
                    )),
                }
            }
        }
    }

    async fn put(&self, key: &str, value: Vec<u8>, options: Option<PutOptions>) -> Result<bool> {
        validate_key(key)?;
        validate_value(&value)?;

        let options = options.unwrap_or_default();
        let (partition_key, row_key) = self.split_key(key);

        let expires_at = options.ttl.map(|d| Utc::now() + d);
        let entity =
            create_table_entity(partition_key.clone(), row_key.clone(), &value, expires_at);

        if options.if_not_exists {
            match self
                .client
                .insert_entity(
                    &self.resource_group_name,
                    &self.account_name,
                    &self.table_name,
                    &entity,
                )
                .await
            {
                Ok(_) => Ok(true),
                Err(e) => {
                    use alien_client_core::ErrorData as CloudErrorData;
                    match e.error.as_ref() {
                        Some(CloudErrorData::RemoteResourceConflict { .. }) => {
                            // Expired rows count as ABSENT, matching the
                            // local provider's atomic takeover: Table
                            // Storage has no server-side TTL, so without
                            // this an expired row (e.g. a dead command
                            // lease) blocks conditional puts forever. The
                            // read + If-Match replace is race-safe: a
                            // concurrent taker changes the ETag and this
                            // update loses with a conflict.
                            self.try_take_over_expired(key, &entity).await
                        }
                        _ => Err(crate::error::map_cloud_client_error(
                            e,
                            format!("Failed to insert entity for key '{}'", key),
                            Some(key.to_string()),
                        )),
                    }
                }
            }
        } else {
            // Insert Or Replace (upsert) - matches Azure REST API terminology
            self.client
                .insert_or_replace_entity(
                    &self.resource_group_name,
                    &self.account_name,
                    &self.table_name,
                    &partition_key,
                    &row_key,
                    &entity,
                )
                .await
                .map_err(|e| {
                    crate::error::map_cloud_client_error(
                        e,
                        format!("Failed to upsert entity for key '{}'", key),
                        Some(key.to_string()),
                    )
                })?;
            Ok(true)
        }
    }

    async fn delete(&self, key: &str) -> Result<()> {
        validate_key(key)?;

        let (partition_key, row_key) = self.split_key(key);

        // Delete entity, ignore if not found
        match self
            .client
            .delete_entity(
                &self.resource_group_name,
                &self.account_name,
                &self.table_name,
                &partition_key,
                &row_key,
                None, // No specific ETag constraint
            )
            .await
        {
            Ok(_) => Ok(()),
            Err(e) => {
                use alien_client_core::ErrorData as CloudErrorData;
                match e.error.as_ref() {
                    Some(CloudErrorData::RemoteResourceNotFound { .. }) => Ok(()), // No error if key doesn't exist
                    _ => Err(crate::error::map_cloud_client_error(
                        e,
                        format!("Failed to delete entity for key '{}'", key),
                        Some(key.to_string()),
                    )),
                }
            }
        }
    }

    async fn exists(&self, key: &str) -> Result<bool> {
        validate_key(key)?;

        let (partition_key, row_key) = self.split_key(key);

        match self
            .client
            .get_entity(
                &self.resource_group_name,
                &self.account_name,
                &self.table_name,
                &partition_key,
                &row_key,
                None,
            )
            .await
        {
            Ok(entity) => {
                // Check TTL expiry
                Ok(!is_entity_expired(&entity))
            }
            Err(e) => {
                use alien_client_core::ErrorData as CloudErrorData;
                match e.error.as_ref() {
                    Some(CloudErrorData::RemoteResourceNotFound { .. }) => Ok(false),
                    _ => Err(crate::error::map_cloud_client_error(
                        e,
                        format!("Failed to check existence of entity for key '{}'", key),
                        Some(key.to_string()),
                    )),
                }
            }
        }
    }

    async fn scan_prefix(
        &self,
        prefix: &str,
        limit: Option<usize>,
        cursor: Option<String>,
    ) -> Result<ScanResult> {
        validate_key(prefix)?;
        let limit = limit.unwrap_or(1000);
        let initial = cursor
            .as_deref()
            .map(|cursor| self.decode_cursor(prefix, cursor))
            .transpose()?
            .unwrap_or(CursorState {
                version: 1,
                prefix: prefix.to_string(),
                current_partition: 0,
                continuation: None,
            });
        if limit == 0 {
            return Ok(ScanResult {
                items: Vec::new(),
                next_cursor: cursor,
            });
        }

        let mut items = Vec::with_capacity(limit);
        let mut partition_id = initial.current_partition;
        let mut continuation = initial.continuation;

        while partition_id < self.num_partitions {
            let partition_key = format!("p{}", partition_id);

            // Build filter with BOTH PartitionKey and RowKey conditions
            // Use a range query approach that's compatible with Azure Table Storage
            let prefix_end = format!("{}~", prefix); // Use tilde as it's after most printable chars
            let filter = format!(
                "(PartitionKey eq '{}') and (RowKey ge '{}') and (RowKey lt '{}')",
                partition_key, prefix, prefix_end
            );

            // Note: We'll do TTL filtering client-side to avoid OData syntax issues
            let filter_with_ttl = filter;

            let query_options = EntityQueryOptions {
                filter: Some(filter_with_ttl),
                select: None,
                top: Some(u32::try_from(limit - items.len()).unwrap_or(u32::MAX)),
                continuation: continuation.clone(),
            };

            let response = self
                .client
                .query_entities(
                    &self.resource_group_name,
                    &self.account_name,
                    &self.table_name,
                    Some(query_options),
                )
                .await
                .map_err(|e| {
                    crate::error::map_cloud_client_error(
                        e,
                        format!("Failed to query entities with prefix '{}'", prefix),
                        Some(prefix.to_string()),
                    )
                })?;

            for entity in response.entities {
                if is_entity_expired(&entity) {
                    continue;
                }

                let key = self.combine_key(&entity.partition_key, &entity.row_key);
                let value = extract_value_from_entity(&entity)?;

                items.push((key, value));
            }

            if items.len() == limit {
                let (next_partition, next_continuation) = match response.continuation {
                    Some(continuation) => (partition_id, Some(continuation)),
                    None if partition_id + 1 < self.num_partitions => (partition_id + 1, None),
                    None => {
                        return Ok(ScanResult {
                            items,
                            next_cursor: None,
                        });
                    }
                };
                return Ok(ScanResult {
                    items,
                    next_cursor: Some(self.encode_cursor(&CursorState {
                        version: 1,
                        prefix: prefix.to_string(),
                        current_partition: next_partition,
                        continuation: next_continuation,
                    })?),
                });
            }

            if let Some(next) = response.continuation {
                continuation = Some(next);
            } else {
                partition_id += 1;
                continuation = None;
            }
        }

        Ok(ScanResult {
            items,
            next_cursor: None,
        })
    }
}