/* 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/. */

//! Multi-store coordinator.

use std::{
    collections::{BTreeMap, BTreeSet},
    fmt::Debug,
    num::NonZeroUsize,
    ops::Bound,
    sync::{Arc, Mutex, OnceLock},
};

// `std::collections::HashMap::extract_if` is unstable [1],
// so we use Hashbrown.
//
// [1]: https://github.com/rust-lang/rust/issues/59618
use hashbrown::hash_map::{Entry, HashMap};

use crate::skv::{
    abort::{AbortController, AbortSignal},
    store::{Store, StorePath},
};

/// Manages access to all persistent stores.
///
/// The coordinator is a singleton that lives for the lifetime of
/// the process. It has two main responsibilities:
///
/// 1. Decoupling the lifecycle of the stores from the
///    lifecycle of the [`crate::skv::interface`] objects.
/// 2. Keeping track of which objects are accessing which stores, and
///    closing stores once they're no longer needed.
///
/// ## Why have a coordinator?
///
/// Interface objects are a little cumbersome to work with in pure Rust:
/// they don't have predictable lifetimes; they can only be held by an
/// [`xpcom::RefPtr`]; and they can unintentionally retain the resources
/// they're managing beyond their intended use.
///
/// The coordinator introduces a layer of indirection to help
/// with these challenges.
///
/// ## Using the coordinator
///
/// When an interface object is instantiated, it either requests a new client
/// from the coordinator, or creates a child of an existing client. The object
/// holds on to the client, and uses it to access any stores it needs
/// during its lifecycle.
///
/// When the object is done accessing its stores, it invalidates its client.
/// The coordinator detaches the invalidated client, and its children,
/// from any open stores. This prevents the object from accessing any stores
/// again, even if the [`xpcom::RefPtr`] holding the object isn't released
/// right away.
///
/// ## Object hierarchies
///
/// Clients can have multiple children, grandchildren, great-grandchildren,
/// and so on; recursively. This is useful for managing hierarchies of
/// interface objects, where the child's lifecycle is the same or shorter
/// than its parent's.
///
/// Check out [`crate::skv::interface::KeyValueService`] and
/// [`crate::skv::interface::KeyValueDatabase`] for an example of
/// hierarchies in action.
#[derive(Debug)]
pub struct Coordinator {
    state: Mutex<CoordinatorState>,
}

impl Coordinator {
    fn new() -> Self {
        Coordinator {
            state: Mutex::new(CoordinatorState {
                clients: BTreeMap::new(),
                stores: HashMap::new(),
            }),
        }
    }

    /// Returns the singleton coordinator.
    pub fn get_or_create() -> &'static Self {
        static COORDINATOR: OnceLock<Coordinator> = OnceLock::new();
        COORDINATOR.get_or_init(|| Coordinator::new())
    }

    /// Creates a new coordinator client.
    pub fn client_with_name(&self, name: &'static str) -> CoordinatorClient<'_> {
        let mut state = self.state.lock().unwrap();

        // Make this client the new last sibling by incrementing
        // the first bud of the current last sibling's key.
        let last_sibling_key = state.clients.last_key_value().map(|(key, _)| key);
        let key = ClientKey::new(
            last_sibling_key
                .map(|key| key.first().increment())
                .unwrap_or(ClientKeyBud::MIN),
        );
        let controller = AbortController::new();
        let signal = controller.signal();
        state
            .clients
            .insert(key.clone(), ActiveClient { controller });

        CoordinatorClient {
            coordinator: self,
            name,
            key,
            signal,
        }
    }
}

/// A client owned by an interface object.
#[derive(Clone, Debug)]
pub struct CoordinatorClient<'a> {
    coordinator: &'a Coordinator,
    name: &'static str,
    key: ClientKey,
    signal: AbortSignal,
}

impl<'a> CoordinatorClient<'a> {
    /// Gets or opens a store for this client.
    ///
    /// Stores connect to the underlying physical database lazily,
    /// so this does _not_ block on I/O.
    pub fn store_for_path(&self, path: StorePath) -> Result<Arc<Store>, CoordinatorError> {
        let mut state = self.coordinator.state.lock().unwrap();
        if !state.clients.contains_key(&self.key) {
            return Err(CoordinatorError::Invalidated(self.name));
        }
        match state.stores.entry(path) {
            Entry::Occupied(mut entry) => Ok(entry.get_mut().attach(self.key.clone())),
            Entry::Vacant(entry) => {
                let store = InUseStore::new(Store::new(entry.key().clone()));
                Ok(entry.insert(store).attach(self.key.clone()))
            }
        }
    }

    /// Creates a child of this client.
    pub fn child_with_name(
        &self,
        name: &'static str,
    ) -> Result<CoordinatorClient<'a>, CoordinatorError> {
        let mut state = self.coordinator.state.lock().unwrap();
        if !state.clients.contains_key(&self.key) {
            return Err(CoordinatorError::Invalidated(self.name));
        }

        let child_key = {
            // Make this child the new last child by incrementing
            // the last bud of the current last child's key.
            let last_child_key = {
                let max_child_key = self.key.clone().appending(ClientKeyBud::MAX);
                let children = state
                    .clients
                    .range((
                        // Include all our descendants, but not ourselves.
                        Bound::Excluded(&self.key),
                        Bound::Included(&max_child_key),
                    ))
                    .map(|(key, _)| key)
                    .filter(|key| {
                        // Narrow down our descendants to immediate children.
                        key.len() <= max_child_key.len()
                    });
                children.last()
            };
            self.key.clone().appending(
                last_child_key
                    .map(|key| key.last().increment())
                    .unwrap_or(ClientKeyBud::MIN),
            )
        };
        let controller = AbortController::new();
        let signal = controller.signal();
        state
            .clients
            .insert(child_key.clone(), ActiveClient { controller });

        Ok(CoordinatorClient {
            coordinator: self.coordinator,
            name,
            key: child_key,
            signal,
        })
    }

    /// Invalidates this client.
    ///
    /// Invalidation is recursive: if the client has descendants
    /// (children, grandchildren, etc.), they'll be invalidated, too.
    ///
    /// If the client or any of its descendants are the last clients
    /// of a store, invalidating the client will also close those stores.
    /// **This can block on disk I/O**, so clients should not be
    /// invalidated on the main thread.
    pub fn invalidate(&self) {
        let (abortable_clients, closeable_stores) = {
            let mut state = self.coordinator.state.lock().unwrap();
            let keys = {
                let max_child_key = self.key.clone().appending(ClientKeyBud::MAX);
                state
                    .clients
                    .range(
                        // Include ourselves and all our descendants.
                        &self.key..=&max_child_key,
                    )
                    .map(|(key, _)| key)
                    .cloned()
                    .collect::<Vec<_>>()
            };
            if keys.is_empty() {
                // Already invalidated; no need to do anything.
                return;
            }
            let abortable_clients = keys
                .iter()
                .map(|key| {
                    // Invariant: `keys` only contains keys that exist
                    // in `clients`.
                    state.clients.remove(key).expect("invariant violation")
                })
                .collect::<Vec<_>>();
            let closeable_stores = state
                .stores
                .extract_if(|_, store| {
                    // Detach ourselves from every store.
                    match store.detach(&keys) {
                        DetachResult::StillInUse => false,
                        // If we detached all clients from this store, also
                        // remove the store from the map, so that we can
                        // close it.
                        DetachResult::Closeable => true,
                    }
                })
                .map(|(_, store)| store)
                .collect::<Vec<_>>();
            (abortable_clients, closeable_stores)
        };

        for client in abortable_clients {
            client.controller.abort();
        }

        for store in closeable_stores {
            // Invariant: `into_inner` always succeeds for closeable stores.
            let store = store.into_inner().expect("invariant violation");
            store.close();
        }
    }

    /// Returns a shared reference to this client's cancellation signal.
    ///
    /// The signal will fire when the client, or any of its ancestors,
    /// are invalidated.
    pub fn signal(&self) -> AbortSignal {
        self.signal.clone()
    }
}

/// An incrementable and totally ordered scalar value.
///
/// These are called "buds" because they can form leaves and branches.
#[derive(Clone, Copy, Eq, Hash, Ord, PartialEq, PartialOrd)]
struct ClientKeyBud(NonZeroUsize);

impl ClientKeyBud {
    const MIN: ClientKeyBud = ClientKeyBud(NonZeroUsize::MIN);
    const MAX: ClientKeyBud = ClientKeyBud(NonZeroUsize::MAX);

    fn increment(self) -> ClientKeyBud {
        Self(self.0.checked_add(1).unwrap())
    }
}

impl Debug for ClientKeyBud {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "#{:?}", self.0)
    }
}

/// A hierarchical key uniquely identifying a coordinator client.
///
/// A key comprises one or more "buds". If clients M and N are siblings,
/// then their keys only differ in the last bud. If N is a child of M,
/// then N's key has one more bud than M's key.
#[derive(Clone, Eq, Hash, Ord, PartialEq, PartialOrd)]
struct ClientKey(ClientKeyBud, Box<[ClientKeyBud]>);

impl ClientKey {
    fn new(id: ClientKeyBud) -> Self {
        Self(id, Box::default())
    }

    /// Returns the first bud of the key.
    fn first(&self) -> ClientKeyBud {
        self.0
    }

    /// Returns the last bud of the key.
    fn last(&self) -> ClientKeyBud {
        self.1.last().copied().unwrap_or(self.0)
    }

    /// Returns the number of buds in the key.
    fn len(&self) -> usize {
        1usize.checked_add(self.1.len()).unwrap()
    }

    /// Consumes this key, and returns a new key with
    /// the given bud added to the end.
    fn appending(self, bud: ClientKeyBud) -> Self {
        let mut buds = self.1.into_vec();
        buds.push(bud);
        ClientKey(self.0, buds.into_boxed_slice())
    }
}

impl Debug for ClientKey {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_set().entry(&self.0).entries(self.1.iter()).finish()
    }
}

#[derive(Debug)]
struct InUseStore {
    attached_client_keys: BTreeSet<ClientKey>,
    store: Arc<Store>,
}

impl InUseStore {
    fn new(store: Store) -> Self {
        Self {
            attached_client_keys: BTreeSet::new(),
            store: Arc::new(store),
        }
    }

    /// Notes that a client with the given key
    /// is now using this store.
    ///
    /// If the client is already using this store,
    /// attaching it again is a no-op.
    fn attach(&mut self, key: ClientKey) -> Arc<Store> {
        self.attached_client_keys.insert(key);
        Arc::clone(&self.store)
    }

    /// Notes that all the clients with the given keys
    /// are no longer using this store.
    fn detach(&mut self, keys: &[ClientKey]) -> DetachResult {
        for key in keys {
            self.attached_client_keys.remove(key);
        }
        if self.attached_client_keys.is_empty() {
            DetachResult::Closeable
        } else {
            DetachResult::StillInUse
        }
    }

    fn into_inner(self) -> Result<Arc<Store>, Self> {
        if self.attached_client_keys.is_empty() {
            Ok(self.store)
        } else {
            Err(self)
        }
    }
}

#[derive(Clone, Copy, Debug)]
enum DetachResult {
    StillInUse,
    Closeable,
}

#[derive(Debug)]
struct ActiveClient {
    controller: AbortController,
}

#[derive(Debug)]
struct CoordinatorState {
    clients: BTreeMap<ClientKey, ActiveClient>,
    stores: HashMap<StorePath, InUseStore>,
}

#[derive(thiserror::Error, Debug)]
pub enum CoordinatorError {
    #[error("`{0}` invalidated")]
    Invalidated(&'static str),
}