// Copyright (c) 2024-present, fjall-rs
// This source code is licensed under both the Apache 2.0 and MIT License
// (found in the LICENSE-* files in the repository)

use crate::{Keyspace, worker_pool::WorkerMessage};
use lsm_tree::{Ingestion as TreeIngestion, UserKey, UserValue};

pub struct Ingestion<'a> {
    keyspace: &'a Keyspace,
    inner: TreeIngestion<'a>,
}

impl<'a> Ingestion<'a> {
    pub fn new(keyspace: &'a Keyspace) -> crate::Result<Self> {
        let inner = TreeIngestion::new(&keyspace.tree)?;
        Ok(Self { keyspace, inner })
    }

    pub fn write<K: Into<UserKey>, V: Into<UserValue>>(
        &mut self,
        key: K,
        value: V,
    ) -> crate::Result<()> {
        self.inner
            .write(key.into(), value.into())
            .map_err(Into::into)
    }

    pub fn write_tombstone<K: Into<UserKey>>(&mut self, key: K) -> crate::Result<()> {
        self.inner.write_tombstone(key.into()).map_err(Into::into)
    }

    #[doc(hidden)]
    pub fn write_weak_tombstone<K: Into<UserKey>>(&mut self, key: K) -> crate::Result<()> {
        self.inner
            .write_weak_tombstone(key.into())
            .map_err(Into::into)
    }

    pub fn finish(self) -> crate::Result<()> {
        // NOTE: We hold to avoid a race condition with concurrent writes:
        //
        // write         ingest
        // lock journal
        // |
        // next seqno=1
        // |
        // --------------finish
        //                 flush
        //                 seqno=2
        //                 register
        //                 |
        // -----------------
        // |
        // insert seqno=1
        let _journal_lock = self.keyspace.supervisor.journal.get_writer();

        self.finish_inner(false)
    }

    /// Finishes the ingestion without taking the global journal writer lock.
    ///
    /// The caller must ensure that no journaled writes can run concurrently in
    /// the same database. Exclusive ingestions into independent keyspaces may
    /// still finish concurrently.
    pub fn finish_exclusive(self) -> crate::Result<()> {
        self.finish_inner(true)
    }

    fn finish_inner(self, exclusive: bool) -> crate::Result<()> {
        let result = if exclusive {
            self.inner.finish_exclusive()
        } else {
            self.inner.finish()
        };

        result
            .inspect(|()| {
                self.keyspace
                    .worker_messager
                    .try_send(WorkerMessage::Compact(self.keyspace.clone()))
                    .ok();

                self.keyspace.supervisor.snapshot_tracker.gc();
            })
            .map_err(Into::into)
    }
}
