// Copyright © 2023 Kim Altintop <kim@eagain.io>
// SPDX-License-Identifier: GPL-2.0-only

use std::{
    path::Path,
    sync::Arc,
};

use automerge::ChangeHash;
use libpijul::{
    pristine::L64,
    small_string::{
        SmallStr,
        SmallString,
    },
};
use sanakirja::{
    btree::{
        self,
        BTreePage,
        Db,
        UDb,
    },
    direct_repr,
    Commit,
    RootDb,
    Storable,
    UnsizedStorable,
};
use thiserror::Error;

#[cfg(test)]
mod tests;

#[derive(Clone, Copy, Debug, PartialEq)]
#[repr(usize)]
pub enum Root {
    Version,
    Docs,
}

const VERSION: L64 = L64(1u64.to_le());

#[derive(Debug, Error)]
pub enum Error {
    #[error("already exists: {0}")]
    AlreadyExists(String),
    #[error("version mismatch")]
    Version,
    #[error("missing root: {0:?}")]
    MissingRoot(Root),
    #[error(transparent)]
    Sanakirja(#[from] sanakirja::Error),
}

pub struct Sanakirja {
    env: Arc<sanakirja::Env>,
}

impl Sanakirja {
    pub fn new(name: Option<&Path>) -> Result<Self, Error> {
        let sz = 1 << 20;
        let roots = 2;
        match name {
            None => sanakirja::Env::new_anon(sz, roots),
            Some(name) => sanakirja::Env::new(name, sz, roots),
        }
        .map(Self::from)
        .map_err(Into::into)
    }

    pub fn txn_begin(&self) -> Result<Txn, Error> {
        let txn = sanakirja::Env::txn_begin(self.env.clone())?;
        if L64(txn.root(Root::Version as usize)) != VERSION {
            return Err(Error::Version);
        }

        let docs = txn
            .root_db(Root::Docs as usize)
            .ok_or_else(|| Error::MissingRoot(Root::Docs))?;

        Ok(Txn { docs, txn })
    }

    pub fn mut_txn_begin(&self) -> Result<MutTxn<()>, Error> {
        let mut txn = sanakirja::Env::mut_txn_begin(self.env.clone())?;
        if let Some(version) = txn.root(Root::Version as usize) {
            if L64(version) != VERSION {
                return Err(Error::Version);
            }
        } else {
            txn.set_root(Root::Version as usize, VERSION.0);
        }

        let docs = txn
            .root_db(Root::Docs as usize)
            .map_or_else(|| btree::create_db_(&mut txn), Ok)?;

        Ok(MutTxn { docs, txn })
    }
}

impl From<sanakirja::Env> for Sanakirja {
    fn from(env: sanakirja::Env) -> Self {
        Self { env: Arc::new(env) }
    }
}

pub struct Doc {
    changes: Db<L64, StoredHash>,
}

impl Doc {
    /// Iterator over the changes in this document, in the order they were
    /// applied (i.e. oldest first).
    pub fn changes<'a, T>(
        &'a self,
        txn: &'a GenericTxn<T>,
    ) -> Changes<'a, T, impl BTreePage<L64, StoredHash>>
    where
        T: sanakirja::LoadPage<Error = sanakirja::Error>,
    {
        Changes {
            inner: Some(btree::iter(&txn.txn, &self.changes, None)),
        }
    }

    /// Iterator over the changes in this document, in reverse order (i.e.
    /// newest first).
    pub fn rev_changes<'a, T>(
        &'a self,
        txn: &'a GenericTxn<T>,
    ) -> RevChanges<'a, T, impl BTreePage<L64, StoredHash>>
    where
        T: sanakirja::LoadPage<Error = sanakirja::Error>,
    {
        RevChanges {
            inner: Some(btree::rev_iter(&txn.txn, &self.changes, None)),
        }
    }

    /// Record a change to the document.
    pub fn put_change<T>(&mut self, txn: &mut MutTxn<T>, change: ChangeHash) -> Result<(), Error> {
        self.extend_with(txn, Some(change))
    }

    /// Record a series of changes to the document.
    pub fn extend_with<T, I>(&mut self, txn: &mut MutTxn<T>, changes: I) -> Result<(), Error>
    where
        I: IntoIterator<Item = ChangeHash>,
    {
        let mut last = if let Some(max) = btree::rev_iter(&txn.txn, &self.changes, None)?.next() {
            let (L64(x), _) = max?;
            *x
        } else {
            0
        };

        for change in changes {
            last += 1;
            btree::put(&mut txn.txn, &mut self.changes, &L64(last), &change.into())?;
        }

        Ok(())
    }
}

/// Iterator over the changes in a [`Doc`] in the order they were applied.
///
/// Created by [`Doc::changes`].
#[must_use = "iterators are lazy and do nothing unless consumed"]
pub struct Changes<'a, T, P>
where
    T: sanakirja::LoadPage,
    P: BTreePage<L64, StoredHash>,
{
    inner: Option<Result<btree::Iter<'a, T, L64, StoredHash, P>, sanakirja::Error>>,
}

impl<'a, T, P> Iterator for Changes<'a, T, P>
where
    T: sanakirja::LoadPage<Error = sanakirja::Error>,
    P: BTreePage<L64, StoredHash> + 'a,
{
    type Item = Result<(u64, ChangeHash), Error>;

    fn next(&mut self) -> Option<Self::Item> {
        let inner = self.inner.take()?;
        inner.map_or_else(
            |e| Some(Err(e.into())),
            |mut iter| {
                iter.next().map(|page| {
                    let x = page
                        .map_err(Into::into)
                        .map(|(t, hash)| (t.0, ChangeHash::from(*hash)));
                    self.inner = Some(Ok(iter));
                    x
                })
            },
        )
    }
}

/// Iterator over the changes in a [`Doc`] in reverse order (newest first).
///
/// Created by [`Doc::rev_iter`].
#[must_use = "iterators are lazy and do nothing unless consumed"]
pub struct RevChanges<'a, T, P>
where
    T: sanakirja::LoadPage,
    P: BTreePage<L64, StoredHash>,
{
    inner: Option<Result<btree::RevIter<'a, T, L64, StoredHash, P>, sanakirja::Error>>,
}

impl<'a, T, P> Iterator for RevChanges<'a, T, P>
where
    T: sanakirja::LoadPage<Error = sanakirja::Error>,
    P: BTreePage<L64, StoredHash> + 'a,
{
    type Item = Result<(u64, ChangeHash), Error>;

    fn next(&mut self) -> Option<Self::Item> {
        let inner = self.inner.take()?;
        inner.map_or_else(
            |e| Some(Err(e.into())),
            |mut iter| {
                iter.next().map(|page| {
                    let x = page
                        .map_err(Into::into)
                        .map(|(t, hash)| (t.0, ChangeHash::from(*hash)));
                    self.inner = Some(Ok(iter));
                    x
                })
            },
        )
    }
}

#[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd)]
#[repr(C)]
pub struct StoredDoc {
    changes: L64,
}

direct_repr!(StoredDoc);
impl sanakirja::debug::Check for StoredDoc {}

// TODO:
// * Versioning?
// * Ask automerge to export HASH_SIZE
#[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd)]
#[repr(C)]
pub struct StoredHash([u8; 32]);

direct_repr!(StoredHash);
impl sanakirja::debug::Check for StoredHash {}

impl From<ChangeHash> for StoredHash {
    fn from(ChangeHash(h): ChangeHash) -> Self {
        Self(h)
    }
}

impl From<StoredHash> for ChangeHash {
    fn from(StoredHash(h): StoredHash) -> Self {
        Self(h)
    }
}

pub type Txn = GenericTxn<sanakirja::Txn<Arc<sanakirja::Env>>>;
pub type MutTxn<T> = GenericTxn<sanakirja::MutTxn<Arc<sanakirja::Env>, T>>;

pub struct GenericTxn<T> {
    txn: T,
    // All documents, by name
    docs: UDb<SmallStr, StoredDoc>,
}

impl<T> GenericTxn<T>
where
    T: sanakirja::LoadPage<Error = sanakirja::Error>,
{
    pub fn get_doc(&self, name: &str) -> Result<Option<Doc>, Error> {
        let name = SmallString::from_str(name);
        let doc = match btree::get(&self.txn, &self.docs, &name, None)? {
            Some((k, v)) if k == name.as_ref() => Some(Doc {
                changes: Db::from_page(v.changes.into()),
            }),

            _ => None,
        };

        Ok(doc)
    }

    pub fn has_doc(&self, name: &str) -> Result<bool, Error> {
        let name = SmallString::from_str(name);
        if let Some((k, _)) = btree::get(&self.txn, &self.docs, &name, None)? {
            Ok(k == name.as_ref())
        } else {
            Ok(false)
        }
    }
}

impl<T> MutTxn<T> {
    pub fn create_doc(&mut self, name: &str) -> Result<Doc, Error> {
        let name = SmallString::from_str(name);
        if let Some((k, _)) = btree::get(&self.txn, &self.docs, &name, None)? {
            if k == name.as_ref() {
                return Err(Error::AlreadyExists(name.as_str().to_owned()));
            }
        }

        let doc = Doc {
            changes: btree::create_db(&mut self.txn)?,
        };

        btree::del(&mut self.txn, &mut self.docs, &name, None)?;
        let stored = StoredDoc {
            changes: doc.changes.db.into(),
        };
        btree::put(&mut self.txn, &mut self.docs, &name, &stored)?;

        Ok(doc)
    }
}

impl Commit for MutTxn<()> {
    fn commit(self) -> Result<(), sanakirja::Error> {
        self.txn.commit()
    }
}