use std::ops::Bound;

use fjall::Readable;

use crate::{DbError, DbResult};

use super::protocol::UpsertKey;

pub(crate) trait RpcHandler: Send + Sync {
    fn get(&self, key: &[u8]) -> DbResult<Option<Vec<u8>>>;
    fn contains(&self, key: &[u8]) -> DbResult<Option<u32>>;
    fn first(&self) -> DbResult<Option<(Vec<u8>, Vec<u8>)>>;
    fn last(&self) -> DbResult<Option<(Vec<u8>, Vec<u8>)>>;
    fn range(
        &self,
        start: Bound<Vec<u8>>,
        end: Bound<Vec<u8>>,
    ) -> DbResult<Vec<(Vec<u8>, Vec<u8>)>>;
    fn range_keys(&self, start: Bound<Vec<u8>>, end: Bound<Vec<u8>>) -> DbResult<Vec<Vec<u8>>>;
    fn upsert(&self, key: UpsertKey, flag: Option<bool>, value: Vec<u8>) -> DbResult<Vec<u8>>;
    fn remove(&self, key: &[u8], soft: bool) -> DbResult<()>;
    fn take(&self, key: &[u8], soft: bool) -> DbResult<Option<Vec<u8>>>;
    fn apply_batch(&self, items: Vec<(Vec<u8>, Option<Vec<u8>>)>) -> DbResult<()>;
}

fn guard_to_kv(guard: fjall::Guard) -> Option<(Vec<u8>, Vec<u8>)> {
    match guard.into_inner() {
        Ok((k, v)) => Some((k.to_vec(), v.to_vec())),
        Err(e) => {
            error!(%e);
            None
        }
    }
}

fn guard_to_key(guard: fjall::Guard) -> Option<Vec<u8>> {
    match guard.key() {
        Ok(k) => Some(k.to_vec()),
        Err(e) => {
            error!(%e);
            None
        }
    }
}

impl RpcHandler for crate::logdb::RawTree {
    fn get(&self, key: &[u8]) -> DbResult<Option<Vec<u8>>> {
        crate::logdb::RawTree::get(self, key)
    }

    fn contains(&self, key: &[u8]) -> DbResult<Option<u32>> {
        match self.tree.get(key)? {
            Some(v) => Ok(Some(v.len() as u32)),
            None => Ok(None),
        }
    }

    fn first(&self) -> DbResult<Option<(Vec<u8>, Vec<u8>)>> {
        match self.tree.first_key_value() {
            Some(guard) => Ok(guard_to_kv(guard)),
            None => Ok(None),
        }
    }

    fn last(&self) -> DbResult<Option<(Vec<u8>, Vec<u8>)>> {
        match self.tree.last_key_value() {
            Some(guard) => Ok(guard_to_kv(guard)),
            None => Ok(None),
        }
    }

    fn range(
        &self,
        start: Bound<Vec<u8>>,
        end: Bound<Vec<u8>>,
    ) -> DbResult<Vec<(Vec<u8>, Vec<u8>)>> {
        Ok(self
            .tree
            .range((start, end))
            .filter_map(guard_to_kv)
            .collect())
    }

    fn range_keys(&self, start: Bound<Vec<u8>>, end: Bound<Vec<u8>>) -> DbResult<Vec<Vec<u8>>> {
        Ok(self
            .tree
            .range((start, end))
            .filter_map(guard_to_key)
            .collect())
    }

    fn upsert(&self, key: UpsertKey, flag: Option<bool>, value: Vec<u8>) -> DbResult<Vec<u8>> {
        let key_bytes = match key {
            UpsertKey::Sequence => {
                let id = self.next_id()?;
                id.to_le_bytes().to_vec()
            }
            UpsertKey::Provided(k) => k,
        };

        if let Some(update_only) = flag {
            let exists = self.tree.contains_key(&key_bytes)?;
            if update_only && !exists {
                return Err(DbError::NotFound);
            }
            if !update_only && exists {
                return Err(DbError::Client("already exists"));
            }
        }

        self.tree.insert(&key_bytes, &value)?;
        self.invalidate_hash(&key_bytes);
        Ok(key_bytes)
    }

    fn remove(&self, key: &[u8], soft: bool) -> DbResult<()> {
        if soft && let Some(val) = self.tree.get(key)? {
            let val_ref: &[u8] = &val;
            self.removed.insert(key, val_ref)?;
        }
        self.tree.remove(key)?;
        self.invalidate_hash(key);
        Ok(())
    }

    fn take(&self, _key: &[u8], _soft: bool) -> DbResult<Option<Vec<u8>>> {
        Err(DbError::NotImplemented)
    }

    fn apply_batch(&self, items: Vec<(Vec<u8>, Option<Vec<u8>>)>) -> DbResult<()> {
        crate::logdb::RawTree::apply_batch(self, items.into_iter().map(|(k, v)| (k.into(), v)))
    }
}

impl RpcHandler for crate::transactional::raw::TxRawTree {
    fn get(&self, key: &[u8]) -> DbResult<Option<Vec<u8>>> {
        crate::transactional::raw::TxRawTree::get(self, key)
    }

    fn contains(&self, key: &[u8]) -> DbResult<Option<u32>> {
        match self.tree.get(key)? {
            Some(v) => Ok(Some(v.len() as u32)),
            None => Ok(None),
        }
    }

    fn first(&self) -> DbResult<Option<(Vec<u8>, Vec<u8>)>> {
        match self.tree.first_key_value() {
            Some(guard) => Ok(guard_to_kv(guard)),
            None => Ok(None),
        }
    }

    fn last(&self) -> DbResult<Option<(Vec<u8>, Vec<u8>)>> {
        match self.tree.last_key_value() {
            Some(guard) => Ok(guard_to_kv(guard)),
            None => Ok(None),
        }
    }

    fn range(
        &self,
        start: Bound<Vec<u8>>,
        end: Bound<Vec<u8>>,
    ) -> DbResult<Vec<(Vec<u8>, Vec<u8>)>> {
        let tx = self.db.db.read_tx();
        Ok(tx
            .range(&self.tree, (start, end))
            .filter_map(guard_to_kv)
            .collect())
    }

    fn range_keys(&self, start: Bound<Vec<u8>>, end: Bound<Vec<u8>>) -> DbResult<Vec<Vec<u8>>> {
        let tx = self.db.db.read_tx();
        Ok(tx
            .range(&self.tree, (start, end))
            .filter_map(guard_to_key)
            .collect())
    }

    fn upsert(&self, key: UpsertKey, flag: Option<bool>, value: Vec<u8>) -> DbResult<Vec<u8>> {
        let key_bytes = match key {
            UpsertKey::Sequence => {
                let id = self.next_id()?;
                id.to_le_bytes().to_vec()
            }
            UpsertKey::Provided(k) => k,
        };

        (|| loop {
            let mut tx = self.db.db.write_tx().map_err(DbError::from)?;

            if let Some(update_only) = flag {
                let exists = tx
                    .contains_key(&self.tree, &key_bytes)
                    .map_err(DbError::from)?;
                if update_only && !exists {
                    return Err(DbError::NotFound);
                }
                if !update_only && exists {
                    return Err(DbError::Client("already exists"));
                }
            }

            tx.insert(&self.tree, &key_bytes, &value);

            match tx.commit() {
                Ok(_) => {
                    self.invalidate_hash(&key_bytes);
                    return Ok(key_bytes.clone());
                }
                Err(_) => continue,
            }
        })()
    }

    fn remove(&self, key: &[u8], soft: bool) -> DbResult<()> {
        (|| loop {
            let mut tx = self.db.db.write_tx().map_err(DbError::from)?;

            if soft && let Some(val) = tx.get(&self.tree, key).map_err(DbError::from)? {
                let val_ref: &[u8] = &val;
                self.removed.insert(key, val_ref)?;
            }

            tx.remove(&self.tree, key);

            match tx.commit() {
                Ok(_) => {
                    self.invalidate_hash(key);
                    return Ok(());
                }
                Err(_) => continue,
            }
        })()
    }

    fn take(&self, key: &[u8], soft: bool) -> DbResult<Option<Vec<u8>>> {
        (|| loop {
            let mut tx = self.db.db.write_tx().map_err(DbError::from)?;
            let old = tx.get(&self.tree, key).map_err(DbError::from)?;

            if old.is_some() {
                if soft && let Some(ref val) = old {
                    let val_ref: &[u8] = val;
                    self.removed.insert(key, val_ref)?;
                }
                tx.remove(&self.tree, key);
            }

            match tx.commit() {
                Ok(_) => {
                    if old.is_some() {
                        self.invalidate_hash(key);
                    }
                    return Ok(old.map(|v: fjall::Slice| v.to_vec()));
                }
                Err(_) => continue,
            }
        })()
    }

    fn apply_batch(&self, items: Vec<(Vec<u8>, Option<Vec<u8>>)>) -> DbResult<()> {
        crate::transactional::raw::TxRawTree::apply_batch(
            self,
            items.into_iter().map(|(k, v)| (k.into(), v)),
        )
    }
}
