/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use crate::crypto;
use crate::utils;
use crate::{datastore_filename, LockstoreError, LockstoreKeystore, SecurityLevel, StoredValue};

use kvstore::skv::store::{Store, StorePath};
use kvstore::skv::{Database, DatabaseError, GetOptions, Key};
use std::path::PathBuf;
use std::sync::Arc;

#[derive(Clone)]
pub struct LockstoreDatastore {
    store: Arc<Store>,
    keystore: Arc<LockstoreKeystore>,
    collection_name: String,
    security_level: SecurityLevel,
}

impl LockstoreDatastore {
    pub fn new(
        dir: PathBuf,
        collection_name: String,
        keystore: Arc<LockstoreKeystore>,
        security_level: SecurityLevel,
    ) -> Result<Self, LockstoreError> {
        keystore.get_dek_internal(&collection_name, security_level)?;
        let data_path = dir.join(datastore_filename(&collection_name));
        Self::init(
            StorePath::OnDisk(data_path),
            collection_name,
            keystore,
            security_level,
        )
    }

    /// Creates an in-memory datastore for the given collection.
    ///
    /// Note: each call allocates a fresh, empty in-memory database. Two
    /// `new_in_memory` instances for the same collection — even with different
    /// security levels sharing the same DEK — do not share data. Cross-security-
    /// level access to the same data requires an on-disk datastore (`new`).
    pub fn new_in_memory(
        collection_name: String,
        keystore: Arc<LockstoreKeystore>,
        security_level: SecurityLevel,
    ) -> Result<Self, LockstoreError> {
        keystore.get_dek_internal(&collection_name, security_level)?;
        Self::init(
            StorePath::for_in_memory(),
            collection_name,
            keystore,
            security_level,
        )
    }

    fn init(
        store_path: StorePath,
        collection_name: String,
        keystore: Arc<LockstoreKeystore>,
        security_level: SecurityLevel,
    ) -> Result<Self, LockstoreError> {
        let store = Arc::new(Store::new(store_path));
        Ok(Self {
            store,
            keystore,
            collection_name,
            security_level,
        })
    }

    pub fn put(&self, entry_name: &str, data: &[u8]) -> Result<(), LockstoreError> {
        let timestamp = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap_or_default()
            .as_secs();

        let stored = StoredValue {
            data: data.to_vec(),
            timestamp,
        };

        let plaintext = serde_json::to_vec(&stored)?;

        let (dek, cipher_suite, _) = self
            .keystore
            .get_dek_internal(&self.collection_name, self.security_level)?;

        let data_to_store = crypto::encrypt_with_key(&plaintext, &dek, cipher_suite)?;

        let full_key = format!("{}::{}", self.collection_name, entry_name);
        let value = utils::bytes_to_value(&data_to_store)?;

        let db = Database::new(&self.store, &self.collection_name);
        let key_obj = Key::from(full_key.as_str());
        db.put(&[(key_obj, Some(value))])?;

        Ok(())
    }

    pub fn get(&self, entry_name: &str) -> Result<Vec<u8>, LockstoreError> {
        let full_key = format!("{}::{}", self.collection_name, entry_name);
        let db = Database::new(&self.store, &self.collection_name);
        let key_obj = Key::from(full_key.as_str());

        let value = db
            .get(&key_obj, &GetOptions::default())?
            .ok_or_else(|| LockstoreError::NotFound(entry_name.to_string()))?;

        let stored_bytes = utils::value_to_bytes(&value)?;

        let (dek, _cipher_suite, _) = self
            .keystore
            .get_dek_internal(&self.collection_name, self.security_level)?;

        let plaintext = crypto::decrypt_with_key(&stored_bytes, &dek)?;

        let stored: StoredValue = serde_json::from_slice(&plaintext)?;
        Ok(stored.data)
    }

    pub fn delete(&self, entry_name: &str) -> Result<(), LockstoreError> {
        let full_key = format!("{}::{}", self.collection_name, entry_name);
        let db = Database::new(&self.store, &self.collection_name);
        let key_obj = Key::from(full_key.as_str());

        if !db.has(&key_obj, &GetOptions::default())? {
            return Err(LockstoreError::NotFound(entry_name.to_string()));
        }

        db.delete(&key_obj)?;
        Ok(())
    }

    pub fn keys(&self) -> Result<Vec<String>, LockstoreError> {
        let reader = self.store.reader()?;
        let prefix = format!("{}::", self.collection_name);

        let entries = reader
            .read(|conn| {
                let mut stmt = conn
                    .prepare(
                        "SELECT data.key FROM data
                         JOIN dbs ON data.db_id = dbs.id
                         WHERE dbs.name = ?1
                         AND data.key LIKE ?2
                         ORDER BY data.key",
                    )
                    .map_err(DatabaseError::from)?;

                let pattern = format!("{}%", prefix);
                let entry_strings: Result<Vec<String>, _> = stmt
                    .query_map([&self.collection_name, &pattern], |row| {
                        let key: String = row.get(0)?;
                        Ok(key.strip_prefix(&prefix).unwrap_or(&key).to_string())
                    })
                    .map_err(DatabaseError::from)?
                    .collect();

                entry_strings.map_err(DatabaseError::from)
            })
            .map_err(LockstoreError::Database)?;

        Ok(entries)
    }

    pub fn close(self) {
        self.store.close();
    }
}