prollytree 0.4.1

/*
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

//! GlueSQL custom storage implementation using ProllyTree.
//!
//! This module implements the GlueSQL `Store`, `StoreMut`, and `Transaction`
//! traits to provide SQL query capabilities over ProllyTree's versioned
//! key-value store.
//!
//! # Async/Sync Bridge
//!
//! GlueSQL requires async trait implementations (`#[async_trait]`), but
//! the underlying [`ThreadSafeGitVersionedKvStore`] is fully synchronous
//! and performs blocking file I/O through the Git object database.
//!
//! To avoid blocking the async executor, all store operations are offloaded
//! to Tokio's blocking thread pool via [`tokio::task::spawn_blocking`].
//! The [`ThreadSafeGitVersionedKvStore`] is `Clone + Send + Sync` (backed
//! by `Arc<Mutex<..>>`), so cloning it into `spawn_blocking` closures is
//! cheap and safe.
//!
//! ```text
//! GlueSQL async query
//!   └─ ProllyStorage (async trait impl)
//!        └─ spawn_blocking
//!             └─ ThreadSafeGitVersionedKvStore (sync, mutex-guarded)
//!                  └─ Git object database (file I/O)
//! ```
//!
//! [`ThreadSafeGitVersionedKvStore`]: crate::git::versioned_store::ThreadSafeGitVersionedKvStore

#[cfg(feature = "sql")]
use std::collections::HashMap;

#[cfg(feature = "sql")]
use async_trait::async_trait;
#[cfg(feature = "sql")]
use futures::stream::iter;
#[cfg(feature = "sql")]
use gluesql_core::{
    data::{Key, Schema},
    error::{Error, Result},
    store::{
        AlterTable, CustomFunction, CustomFunctionMut, DataRow, Index, IndexMut, Metadata, Planner,
        RowIter, Store, StoreMut, Transaction,
    },
};

#[cfg(feature = "sql")]
use crate::git::versioned_store::ThreadSafeGitVersionedKvStore;

#[cfg(feature = "sql")]
const LEGACY_SCHEMA_ID: &str = "__schema__";
#[cfg(feature = "sql")]
const ROW_KEY_PREFIX: &[u8] = b":row:";
#[cfg(feature = "sql")]
const SCHEMA_KEY_SUFFIX: &[u8] = b":\0schema";
#[cfg(feature = "sql")]
const LEGACY_SCHEMA_KEY_SUFFIX: &[u8] = b":__schema__";

/// GlueSQL storage backend using ProllyTree.
///
/// Wraps a [`ThreadSafeGitVersionedKvStore`] and implements GlueSQL's async
/// storage traits. All blocking store operations are offloaded to
/// [`tokio::task::spawn_blocking`] to keep the async executor responsive.
///
/// [`ThreadSafeGitVersionedKvStore`]: crate::git::versioned_store::ThreadSafeGitVersionedKvStore
#[cfg(feature = "sql")]
pub struct ProllyStorage<const D: usize> {
    store: ThreadSafeGitVersionedKvStore<D>,
    schemas: HashMap<String, Schema>,
}

#[cfg(feature = "sql")]
impl<const D: usize> ProllyStorage<D> {
    /// Create a new ProllyStorage instance
    pub fn new(store: ThreadSafeGitVersionedKvStore<D>) -> Self {
        Self {
            store,
            schemas: HashMap::new(),
        }
    }

    /// Initialize with a path
    #[allow(clippy::result_large_err)]
    pub fn init(path: &std::path::Path) -> Result<Self> {
        let dir = path.to_path_buf();
        let dir_string = dir.to_string_lossy().to_string();
        let store = ThreadSafeGitVersionedKvStore::init(path).map_err(|e| {
            Error::StorageMsg(format!("Failed to initialize store: {e} from {dir_string}"))
        })?;
        Ok(Self::new(store))
    }

    /// Open an existing storage
    #[allow(clippy::result_large_err)]
    pub fn open(path: &std::path::Path) -> Result<Self> {
        let store = ThreadSafeGitVersionedKvStore::open(path)
            .map_err(|e| Error::StorageMsg(format!("Failed to open store: {e}")))?;
        Ok(Self::new(store))
    }

    // returns the underlying store
    pub fn store(&self) -> &ThreadSafeGitVersionedKvStore<D> {
        &self.store
    }

    /// Convert table name and row key to storage key
    fn make_storage_key(table_name: &str, key: &Key) -> Vec<u8> {
        let mut storage_key = Self::row_prefix(table_name);
        let encoded_key = serde_json::to_vec(key).expect("GlueSQL key serialization cannot fail");
        storage_key.extend_from_slice(&encoded_key);
        storage_key
    }

    /// Convert table name and row key to the legacy untyped storage key.
    fn legacy_storage_key(table_name: &str, key: &Key) -> Vec<u8> {
        match key {
            Key::I64(id) => format!("{table_name}:{id}").into_bytes(),
            Key::Str(id) => format!("{table_name}:{id}").into_bytes(),
            Key::None => format!("{table_name}:{LEGACY_SCHEMA_ID}").into_bytes(),
            _ => format!("{table_name}:{key:?}").into_bytes(),
        }
    }

    /// Get schema key for a table
    fn schema_key(table_name: &str) -> Vec<u8> {
        let mut key = table_name.as_bytes().to_vec();
        key.extend_from_slice(SCHEMA_KEY_SUFFIX);
        key
    }

    fn legacy_schema_key(table_name: &str) -> Vec<u8> {
        format!("{table_name}:{LEGACY_SCHEMA_ID}").into_bytes()
    }

    /// Get the typed row key prefix for a table.
    fn row_prefix(table_name: &str) -> Vec<u8> {
        let mut prefix = table_name.as_bytes().to_vec();
        prefix.extend_from_slice(ROW_KEY_PREFIX);
        prefix
    }

    /// Get the legacy table prefix for untyped row keys.
    fn legacy_table_prefix(table_name: &str) -> Vec<u8> {
        format!("{table_name}:").into_bytes()
    }

    /// Parse key from storage key string
    fn parse_key_from_storage_key(storage_key: &[u8], table_prefix: &str) -> Result<Key> {
        let row_prefix = Self::row_prefix(table_prefix);
        if let Some(encoded_key) = storage_key.strip_prefix(row_prefix.as_slice()) {
            let key = serde_json::from_slice(encoded_key).map_err(|e| {
                Error::StorageMsg(format!("Failed to deserialize typed row key: {e}"))
            })?;
            return Ok(key);
        }

        Ok(Self::parse_legacy_key_from_storage_key(
            storage_key,
            table_prefix,
        ))
    }

    /// Parse a key from the legacy untyped storage key string.
    fn parse_legacy_key_from_storage_key(storage_key: &[u8], table_prefix: &str) -> Key {
        let key_str = String::from_utf8_lossy(storage_key);
        let prefix = format!("{table_prefix}:");
        let key_part = key_str.strip_prefix(&prefix).unwrap_or("");

        if let Ok(id) = key_part.parse::<i64>() {
            Key::I64(id)
        } else {
            Key::Str(key_part.to_string())
        }
    }

    /// Check whether a key is the legacy schema sentinel for this table.
    fn is_schema_storage_key(table_name: &str, storage_key: &[u8]) -> bool {
        storage_key == Self::schema_key(table_name).as_slice()
            || storage_key == Self::legacy_schema_key(table_name).as_slice()
    }

    /// Commit with a custom message.
    ///
    /// Uses `spawn_blocking` to offload the synchronous git commit to a
    /// blocking thread, keeping the async executor free.
    pub async fn commit_with_message(&mut self, message: &str) -> Result<()> {
        let store = self.store.clone();
        let message = message.to_string();
        tokio::task::spawn_blocking(move || {
            store
                .commit(&message)
                .map_err(|e| Error::StorageMsg(format!("Failed to commit: {e}")))
        })
        .await
        .map_err(|e| Error::StorageMsg(format!("Blocking task join error: {e}")))??;
        Ok(())
    }
}

// Implement all the required traits
#[cfg(feature = "sql")]
impl<const D: usize> AlterTable for ProllyStorage<D> {}
#[cfg(feature = "sql")]
impl<const D: usize> Index for ProllyStorage<D> {}
#[cfg(feature = "sql")]
impl<const D: usize> IndexMut for ProllyStorage<D> {}
#[cfg(feature = "sql")]
impl<const D: usize> Metadata for ProllyStorage<D> {}
#[cfg(feature = "sql")]
impl<const D: usize> CustomFunction for ProllyStorage<D> {}
#[cfg(feature = "sql")]
impl<const D: usize> CustomFunctionMut for ProllyStorage<D> {}
#[cfg(feature = "sql")]
impl<const D: usize> Planner for ProllyStorage<D> {}

#[cfg(feature = "sql")]
#[async_trait]
impl<const D: usize> Store for ProllyStorage<D> {
    async fn fetch_all_schemas(&self) -> Result<Vec<Schema>> {
        let store = self.store.clone();
        tokio::task::spawn_blocking(move || {
            let all_keys = store
                .list_keys()
                .map_err(|e| Error::StorageMsg(format!("Failed to list keys: {e}")))?;
            let mut schemas = HashMap::new();

            for storage_key in all_keys {
                let table_name_bytes = if let Some(table_name) =
                    storage_key.strip_suffix(SCHEMA_KEY_SUFFIX)
                {
                    table_name
                } else if let Some(table_name) = storage_key.strip_suffix(LEGACY_SCHEMA_KEY_SUFFIX)
                {
                    table_name
                } else {
                    continue;
                };
                let table_name = String::from_utf8_lossy(table_name_bytes);
                if Self::is_schema_storage_key(&table_name, &storage_key) {
                    if let Some(schema_data) = store.get(&storage_key) {
                        let schema: Schema = serde_json::from_slice(&schema_data).map_err(|e| {
                            Error::StorageMsg(format!("Failed to deserialize schema: {e}"))
                        })?;
                        schemas.insert(schema.table_name.clone(), schema);
                    }
                }
            }

            let mut schemas: Vec<_> = schemas.into_values().collect();
            schemas.sort_by(|a, b| a.table_name.cmp(&b.table_name));
            Ok(schemas)
        })
        .await
        .map_err(|e| Error::StorageMsg(format!("Blocking task join error: {e}")))?
    }

    async fn fetch_schema(&self, table_name: &str) -> Result<Option<Schema>> {
        let store = self.store.clone();
        let key = Self::schema_key(table_name);
        let legacy_key = Self::legacy_schema_key(table_name);
        tokio::task::spawn_blocking(move || {
            if let Some(schema_data) = store.get(&key).or_else(|| store.get(&legacy_key)) {
                let schema: Schema = serde_json::from_slice(&schema_data)
                    .map_err(|e| Error::StorageMsg(format!("Failed to deserialize schema: {e}")))?;
                Ok(Some(schema))
            } else {
                Ok(None)
            }
        })
        .await
        .map_err(|e| Error::StorageMsg(format!("Blocking task join error: {e}")))?
    }

    async fn fetch_data(&self, table_name: &str, key: &Key) -> Result<Option<DataRow>> {
        let store = self.store.clone();
        let storage_key = Self::make_storage_key(table_name, key);
        let legacy_storage_key = Self::legacy_storage_key(table_name, key);
        let legacy_schema_key = Self::legacy_schema_key(table_name);
        tokio::task::spawn_blocking(move || {
            let row_data = store.get(&storage_key).or_else(|| {
                if legacy_storage_key == legacy_schema_key {
                    None
                } else {
                    store.get(&legacy_storage_key)
                }
            });
            if let Some(row_data) = row_data {
                let row: DataRow = serde_json::from_slice(&row_data)
                    .map_err(|e| Error::StorageMsg(format!("Failed to deserialize row: {e}")))?;
                Ok(Some(row))
            } else {
                Ok(None)
            }
        })
        .await
        .map_err(|e| Error::StorageMsg(format!("Blocking task join error: {e}")))?
    }

    async fn scan_data<'a>(&'a self, table_name: &str) -> Result<RowIter> {
        let store = self.store.clone();
        let table_name = table_name.to_string();
        tokio::task::spawn_blocking(move || {
            let row_prefix = Self::row_prefix(&table_name);
            let legacy_prefix = Self::legacy_table_prefix(&table_name);

            let all_keys = store
                .list_keys()
                .map_err(|e| Error::StorageMsg(format!("Failed to list keys: {e}")))?;
            let mut rows_by_key = HashMap::new();

            for storage_key in all_keys.iter() {
                if storage_key.starts_with(&legacy_prefix)
                    && !storage_key.starts_with(&row_prefix)
                    && !Self::is_schema_storage_key(&table_name, storage_key)
                {
                    if let Some(row_data) = store.get(storage_key) {
                        let row: DataRow = serde_json::from_slice(&row_data).map_err(|e| {
                            Error::StorageMsg(format!("Failed to deserialize row: {e}"))
                        })?;

                        let key = ProllyStorage::<D>::parse_key_from_storage_key(
                            storage_key,
                            &table_name,
                        )?;
                        rows_by_key.entry(key).or_insert(row);
                    }
                }
            }

            for storage_key in all_keys {
                if storage_key.starts_with(&row_prefix) {
                    if let Some(row_data) = store.get(&storage_key) {
                        let row: DataRow = serde_json::from_slice(&row_data).map_err(|e| {
                            Error::StorageMsg(format!("Failed to deserialize row: {e}"))
                        })?;

                        let key = ProllyStorage::<D>::parse_key_from_storage_key(
                            &storage_key,
                            &table_name,
                        )?;
                        rows_by_key.insert(key, row);
                    }
                }
            }

            let mut rows: Vec<_> = rows_by_key
                .into_iter()
                .map(|(key, row)| Ok((key, row)))
                .collect();

            rows.sort_by(|a, b| match (a, b) {
                (Ok((key_a, _)), Ok((key_b, _))) => key_a.cmp(key_b),
                _ => std::cmp::Ordering::Equal,
            });

            Ok(Box::pin(iter(rows)) as RowIter)
        })
        .await
        .map_err(|e| Error::StorageMsg(format!("Blocking task join error: {e}")))?
    }
}

#[cfg(feature = "sql")]
#[async_trait]
impl<const D: usize> StoreMut for ProllyStorage<D> {
    async fn insert_schema(&mut self, schema: &Schema) -> Result<()> {
        let store = self.store.clone();
        let key = Self::schema_key(&schema.table_name);
        let legacy_key = Self::legacy_schema_key(&schema.table_name);
        let schema_data = serde_json::to_vec(schema)
            .map_err(|e| Error::StorageMsg(format!("Failed to serialize schema: {e}")))?;

        tokio::task::spawn_blocking(move || {
            store
                .insert(key, schema_data)
                .map_err(|e| Error::StorageMsg(format!("Failed to insert schema: {e}")))?;
            store.delete(&legacy_key).map_err(|e| {
                Error::StorageMsg(format!("Failed to delete legacy schema key: {e}"))
            })?;
            Ok::<(), Error>(())
        })
        .await
        .map_err(|e| Error::StorageMsg(format!("Blocking task join error: {e}")))??;

        // Cache the schema (back on the async task, no I/O)
        self.schemas
            .insert(schema.table_name.clone(), schema.clone());

        Ok(())
    }

    async fn delete_schema(&mut self, table_name: &str) -> Result<()> {
        let store = self.store.clone();
        let key = Self::schema_key(table_name);
        let legacy_key = Self::legacy_schema_key(table_name);

        tokio::task::spawn_blocking(move || {
            store
                .delete(&key)
                .map_err(|e| Error::StorageMsg(format!("Failed to delete schema: {e}")))?;
            store.delete(&legacy_key).map_err(|e| {
                Error::StorageMsg(format!("Failed to delete legacy schema key: {e}"))
            })?;
            Ok::<(), Error>(())
        })
        .await
        .map_err(|e| Error::StorageMsg(format!("Blocking task join error: {e}")))??;

        // Remove from cache
        self.schemas.remove(table_name);

        Ok(())
    }

    async fn append_data(&mut self, table_name: &str, rows: Vec<DataRow>) -> Result<()> {
        let store = self.store.clone();
        let table_name = table_name.to_string();
        tokio::task::spawn_blocking(move || {
            for row in rows {
                let mut counter = 0i64;
                let storage_key = loop {
                    let key = Key::I64(counter);
                    let storage_key = ProllyStorage::<D>::make_storage_key(&table_name, &key);

                    if store.get(&storage_key).is_none() {
                        break storage_key;
                    }
                    counter += 1;
                };

                let row_data = serde_json::to_vec(&row)
                    .map_err(|e| Error::StorageMsg(format!("Failed to serialize row: {e}")))?;

                store
                    .insert(storage_key, row_data)
                    .map_err(|e| Error::StorageMsg(format!("Failed to insert row: {e}")))?;
            }
            Ok(())
        })
        .await
        .map_err(|e| Error::StorageMsg(format!("Blocking task join error: {e}")))?
    }

    async fn insert_data(&mut self, table_name: &str, rows: Vec<(Key, DataRow)>) -> Result<()> {
        let store = self.store.clone();
        let table_name = table_name.to_string();
        tokio::task::spawn_blocking(move || {
            for (key, row) in rows {
                let storage_key = ProllyStorage::<D>::make_storage_key(&table_name, &key);
                let row_data = serde_json::to_vec(&row)
                    .map_err(|e| Error::StorageMsg(format!("Failed to serialize row: {e}")))?;

                store
                    .insert(storage_key, row_data)
                    .map_err(|e| Error::StorageMsg(format!("Failed to insert row: {e}")))?;
            }
            Ok(())
        })
        .await
        .map_err(|e| Error::StorageMsg(format!("Blocking task join error: {e}")))?
    }

    async fn delete_data(&mut self, table_name: &str, keys: Vec<Key>) -> Result<()> {
        let store = self.store.clone();
        let table_name = table_name.to_string();
        tokio::task::spawn_blocking(move || {
            for key in keys {
                let storage_key = ProllyStorage::<D>::make_storage_key(&table_name, &key);
                let legacy_storage_key = ProllyStorage::<D>::legacy_storage_key(&table_name, &key);
                let legacy_schema_key = ProllyStorage::<D>::legacy_schema_key(&table_name);

                store
                    .delete(&storage_key)
                    .map_err(|e| Error::StorageMsg(format!("Failed to delete row: {e}")))?;
                if legacy_storage_key != legacy_schema_key {
                    store.delete(&legacy_storage_key).map_err(|e| {
                        Error::StorageMsg(format!("Failed to delete legacy row: {e}"))
                    })?;
                }
            }
            Ok(())
        })
        .await
        .map_err(|e| Error::StorageMsg(format!("Blocking task join error: {e}")))?
    }
}

#[cfg(feature = "sql")]
#[async_trait]
impl<const D: usize> Transaction for ProllyStorage<D> {
    async fn begin(&mut self, autocommit: bool) -> Result<bool> {
        if autocommit {
            return Ok(false);
        }

        // ProllyTree with git backend doesn't support nested transactions.
        // Always return false to indicate no transaction was started.
        Ok(false)
    }

    async fn rollback(&mut self) -> Result<()> {
        // Since we don't support transactions, rollback is a no-op.
        Ok(())
    }

    async fn commit(&mut self) -> Result<()> {
        let store = self.store.clone();
        tokio::task::spawn_blocking(move || {
            store
                .commit("Transaction commit")
                .map_err(|e| Error::StorageMsg(format!("Failed to commit transaction: {e}")))
        })
        .await
        .map_err(|e| Error::StorageMsg(format!("Blocking task join error: {e}")))??;
        Ok(())
    }
}

#[cfg(all(test, feature = "sql"))]
mod tests {
    use super::*;
    use gluesql_core::{
        ast::{ColumnDef, DataType},
        data::{Key, Schema, Value},
        store::DataRow,
    };
    use tempfile::TempDir;

    #[tokio::test]
    async fn test_basic_operations() {
        let temp_dir = TempDir::new().unwrap();

        // Initialize git repository in temp directory
        std::process::Command::new("git")
            .arg("init")
            .current_dir(temp_dir.path())
            .output()
            .expect("Failed to initialize git repository");

        // Create a subdirectory for the dataset
        let dataset_path = temp_dir.path().join("dataset");
        std::fs::create_dir(&dataset_path).unwrap();

        let mut storage = ProllyStorage::<32>::init(&dataset_path).unwrap();

        // Create a simple schema
        let schema = Schema {
            table_name: "users".to_string(),
            column_defs: Some(vec![
                ColumnDef {
                    name: "id".to_string(),
                    data_type: DataType::Int,
                    nullable: false,
                    default: None,
                    unique: None,
                    comment: None,
                },
                ColumnDef {
                    name: "name".to_string(),
                    data_type: DataType::Text,
                    nullable: false,
                    default: None,
                    unique: None,
                    comment: None,
                },
            ]),
            indexes: vec![],
            engine: None,
            foreign_keys: vec![],
            comment: None,
        };

        // Insert schema
        storage.insert_schema(&schema).await.unwrap();

        // Verify schema
        let fetched = storage.fetch_schema("users").await.unwrap();
        assert!(fetched.is_some());

        // Insert some data
        let row = DataRow::Vec(vec![Value::I64(1), Value::Str("Alice".to_string())]);
        let key = Key::I64(1);
        storage
            .insert_data("users", vec![(key.clone(), row.clone())])
            .await
            .unwrap();

        // Fetch data
        let fetched_row = storage.fetch_data("users", &key).await.unwrap();
        assert!(fetched_row.is_some());

        // Scan data
        use futures::StreamExt;
        let mut iter = storage.scan_data("users").await.unwrap();
        let first = iter.next().await.unwrap().unwrap();
        assert_eq!(first.0, key);
    }

    #[test]
    fn test_text_key_round_trip_preserves_key_variant() {
        let storage_key = ProllyStorage::<32>::make_storage_key("users", &Key::Str("42".into()));
        let parsed =
            ProllyStorage::<32>::parse_key_from_storage_key(&storage_key, "users").unwrap();

        assert_eq!(parsed, Key::Str("42".into()));
    }

    #[test]
    fn test_non_string_key_round_trip_preserves_key_variant() {
        for key in [Key::I64(42), Key::Bool(true), Key::None] {
            let storage_key = ProllyStorage::<32>::make_storage_key("users", &key);
            let parsed =
                ProllyStorage::<32>::parse_key_from_storage_key(&storage_key, "users").unwrap();

            assert_eq!(parsed, key);
        }
    }
}