//! redb-backed JMT node store.
//!
//! [`RedbNodeStore`] owns the redb database. [`TransactionReader`] implements [`TreeReader`]
//! against a single read transaction so JMT traversals don't open a fresh transaction per node.
//! A single unified JMT shares one set of tables — namespace separation is at the key level
//! via `leaf_key[0]`.

use crate::error::{Result, StateTreeError};
use borsh::BorshDeserialize;
use jmt::{
    storage::{LeafNode, Node, NodeKey, TreeReader, TreeUpdateBatch},
    KeyHash, OwnedValue, Version,
};
use redb::{ReadableDatabase, ReadableTable, TableDefinition};
use std::path::Path;

const NODES_TABLE: TableDefinition<'_, &[u8], &[u8]> = TableDefinition::new("nodes");
const VALUES_TABLE: TableDefinition<'_, &[u8], &[u8]> = TableDefinition::new("values");
const TOMBSTONES_TABLE: TableDefinition<'_, &[u8], &[u8]> = TableDefinition::new("tombstones");
const METADATA_TABLE: TableDefinition<'_, &str, &[u8]> = TableDefinition::new("metadata");

const JMT_VERSION_KEY: &str = "jmt_version";
const RIGHTMOST_LEAF_KEY: &str = "rightmost_leaf";

/// redb-backed storage for JMT nodes and values in the unified state tree.
pub(crate) struct RedbNodeStore {
    db: redb::Database,
}

impl std::fmt::Debug for RedbNodeStore {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("RedbNodeStore").finish_non_exhaustive()
    }
}

impl RedbNodeStore {
    /// Open or create the redb database, initializing tables.
    pub(crate) fn open(path: impl AsRef<Path>) -> Result<Self> {
        let db = redb::Database::create(path)?;
        let txn = db.begin_write()?;
        txn.open_table(NODES_TABLE)?;
        txn.open_table(VALUES_TABLE)?;
        txn.open_table(TOMBSTONES_TABLE)?;
        txn.open_table(METADATA_TABLE)?;
        txn.commit()?;
        Ok(Self { db })
    }

    /// Execute a closure within a read transaction.
    pub(crate) fn read<F, R>(&self, f: F) -> Result<R>
    where
        F: FnOnce(&redb::ReadTransaction) -> Result<R>,
    {
        let txn = self.db.begin_read()?;
        f(&txn)
    }

    /// Begin a read transaction (caller manages lifetime).
    pub(crate) fn begin_read(&self) -> Result<redb::ReadTransaction> {
        Ok(self.db.begin_read()?)
    }

    /// Begin a write transaction with `Durability::Immediate` (caller manages commit).
    pub(crate) fn begin_write(&self) -> Result<redb::WriteTransaction> {
        let mut txn = self.db.begin_write()?;
        txn.set_durability(redb::Durability::Immediate)?;
        Ok(txn)
    }
}

/// Read-only view of the unified JMT tables pinned to a single redb read transaction.
pub(crate) struct TransactionReader {
    nodes_table: redb::ReadOnlyTable<&'static [u8], &'static [u8]>,
    values_table: redb::ReadOnlyTable<&'static [u8], &'static [u8]>,
    tombstones_table: redb::ReadOnlyTable<&'static [u8], &'static [u8]>,
    metadata_table: redb::ReadOnlyTable<&'static str, &'static [u8]>,
    #[allow(dead_code)]
    txn: redb::ReadTransaction,
}

impl TransactionReader {
    pub(crate) fn new(txn: redb::ReadTransaction) -> Result<Self> {
        let nodes_table = txn.open_table(NODES_TABLE)?;
        let values_table = txn.open_table(VALUES_TABLE)?;
        let tombstones_table = txn.open_table(TOMBSTONES_TABLE)?;
        let metadata_table = txn.open_table(METADATA_TABLE)?;
        Ok(Self {
            nodes_table,
            values_table,
            tombstones_table,
            metadata_table,
            txn,
        })
    }
}

/// On version tie (value and tombstone at the same version), tombstone wins.
fn resolve_value(
    values: &redb::ReadOnlyTable<&'static [u8], &'static [u8]>,
    tombstones: &redb::ReadOnlyTable<&'static [u8], &'static [u8]>,
    max_version: Version,
    key_hash: KeyHash,
) -> Result<Option<OwnedValue>> {
    let (start, end) = version_range_keys(&key_hash, max_version);

    let tombstone_version = {
        let mut range = tombstones.range::<&[u8]>(start.as_slice()..=end.as_slice())?;
        match range.next_back() {
            Some(entry) => {
                let (k, _) = entry?;
                Some(version_from_key(k.value())?)
            }
            None => None,
        }
    };
    let value_entry = {
        let mut range = values.range::<&[u8]>(start.as_slice()..=end.as_slice())?;
        range.next_back().transpose()?
    };

    let value_version = value_entry
        .as_ref()
        .map(|(k, _)| version_from_key(k.value()))
        .transpose()?;

    match (value_entry, value_version, tombstone_version) {
        (Some((_, v)), Some(vv), Some(tv)) if vv > tv => Ok(Some(v.value().to_vec())),
        (Some((_, v)), Some(_), None) => Ok(Some(v.value().to_vec())),
        _ => Ok(None),
    }
}

impl TreeReader for TransactionReader {
    fn get_node_option(&self, node_key: &NodeKey) -> anyhow::Result<Option<Node>> {
        let nk_bytes = borsh::to_vec(node_key)?;
        match self.nodes_table.get(nk_bytes.as_slice())? {
            Some(guard) => {
                let node = Node::try_from_slice(guard.value())?;
                Ok(Some(node))
            }
            None => Ok(None),
        }
    }

    fn get_value_option(&self, max_version: Version, key_hash: KeyHash) -> anyhow::Result<Option<OwnedValue>> {
        Ok(resolve_value(
            &self.values_table,
            &self.tombstones_table,
            max_version,
            key_hash,
        )?)
    }

    fn get_rightmost_leaf(&self) -> anyhow::Result<Option<(NodeKey, LeafNode)>> {
        match self.metadata_table.get(RIGHTMOST_LEAF_KEY)? {
            Some(guard) => {
                let bytes = guard.value();
                let (node_key, leaf_node) = deserialize_rightmost_leaf(bytes)?;
                Ok(Some((node_key, leaf_node)))
            }
            None => Ok(None),
        }
    }
}

/// Increment the unified JMT version counter, returning the new version.
pub(crate) fn next_version(txn: &redb::WriteTransaction) -> Result<Version> {
    let mut table = txn.open_table(METADATA_TABLE)?;
    let current = {
        let maybe = table.get(JMT_VERSION_KEY)?;
        match maybe {
            Some(guard) => {
                let bytes: [u8; 8] = guard
                    .value()
                    .try_into()
                    .map_err(|_| StateTreeError::Serialization("invalid version bytes in metadata".into()))?;
                Some(u64::from_le_bytes(bytes))
            }
            None => None,
        }
    };
    let version = match current {
        Some(v) => v + 1,
        None => 0,
    };
    table.insert(JMT_VERSION_KEY, version.to_le_bytes().as_slice())?;
    Ok(version)
}

/// Write JMT internal nodes from the update batch, remove stale nodes, and update the
/// rightmost leaf metadata entry.
///
/// The rightmost leaf is advanced when the batch contains a leaf with a larger key_hash.
/// If the stored rightmost's NodeKey appears in the stale set (i.e. the rightmost leaf
/// was deleted or restructured), the metadata is cleared so `get_rightmost_leaf` returns
/// `None`. The restore path treats `None` as a fresh tree and re-ingests from the start.
pub(crate) fn write_nodes(txn: &redb::WriteTransaction, batch: &TreeUpdateBatch) -> Result<()> {
    let mut table = txn.open_table(NODES_TABLE)?;
    let mut new_rightmost: Option<(NodeKey, LeafNode)> = None;

    for (node_key, node) in batch.node_batch.nodes() {
        let nk_bytes = borsh::to_vec(node_key).map_err(|e| StateTreeError::Serialization(e.to_string()))?;
        let n_bytes = borsh::to_vec(node).map_err(|e| StateTreeError::Serialization(e.to_string()))?;
        table.insert(nk_bytes.as_slice(), n_bytes.as_slice())?;

        if let Node::Leaf(leaf_node) = node {
            let is_greater = new_rightmost
                .as_ref()
                .is_none_or(|(_, cur)| leaf_node.key_hash() > cur.key_hash());
            if is_greater {
                new_rightmost = Some((node_key.clone(), leaf_node.clone()));
            }
        }
    }

    let stale_node_keys: std::collections::HashSet<Vec<u8>> = batch
        .stale_node_index_batch
        .iter()
        .map(|s| borsh::to_vec(&s.node_key).map_err(|e| StateTreeError::Serialization(e.to_string())))
        .collect::<Result<_>>()?;

    for nk_bytes in &stale_node_keys {
        table.remove(nk_bytes.as_slice())?;
    }

    let mut meta = txn.open_table(METADATA_TABLE)?;
    let existing = match meta.get(RIGHTMOST_LEAF_KEY)? {
        Some(guard) => {
            let (nk, leaf) =
                deserialize_rightmost_leaf(guard.value()).map_err(|e| StateTreeError::Serialization(e.to_string()))?;
            Some((nk, leaf))
        }
        None => None,
    };

    let existing_was_invalidated = existing
        .as_ref()
        .is_some_and(|(nk, _)| borsh::to_vec(nk).ok().is_some_and(|b| stale_node_keys.contains(&b)));

    if existing_was_invalidated {
        meta.remove(RIGHTMOST_LEAF_KEY)?;
    }

    if let Some((ref nk, ref leaf)) = new_rightmost {
        let should_update = match meta.get(RIGHTMOST_LEAF_KEY)? {
            Some(guard) => {
                let (_, cur) = deserialize_rightmost_leaf(guard.value())
                    .map_err(|e| StateTreeError::Serialization(e.to_string()))?;
                leaf.key_hash() > cur.key_hash()
            }
            None => true,
        };
        if should_update {
            let bytes = serialize_rightmost_leaf(nk, leaf)?;
            meta.insert(RIGHTMOST_LEAF_KEY, bytes.as_slice())?;
        }
    }

    Ok(())
}

#[derive(borsh::BorshSerialize, borsh::BorshDeserialize)]
struct RightmostLeafEntry {
    node_key: NodeKey,
    leaf_node: LeafNode,
}

fn serialize_rightmost_leaf(node_key: &NodeKey, leaf_node: &LeafNode) -> Result<Vec<u8>> {
    let entry = RightmostLeafEntry {
        node_key: node_key.clone(),
        leaf_node: leaf_node.clone(),
    };
    borsh::to_vec(&entry).map_err(|e| StateTreeError::Serialization(e.to_string()))
}

fn deserialize_rightmost_leaf(bytes: &[u8]) -> anyhow::Result<(NodeKey, LeafNode)> {
    let entry = RightmostLeafEntry::try_from_slice(bytes)?;
    Ok((entry.node_key, entry.leaf_node))
}

/// Write a versioned value entry. `Some` writes raw bytes to the values table;
/// `None` inserts a marker in the tombstones table.
pub(crate) fn write_value(
    txn: &redb::WriteTransaction,
    version: Version,
    key_hash: KeyHash,
    value: Option<Vec<u8>>,
) -> Result<()> {
    let vk = version_key(&key_hash, version);
    match value {
        Some(v) => {
            let mut table = txn.open_table(VALUES_TABLE)?;
            table.insert(vk.as_slice(), v.as_slice())?;
        }
        None => {
            let mut table = txn.open_table(TOMBSTONES_TABLE)?;
            table.insert(vk.as_slice(), [].as_slice())?;
        }
    }
    Ok(())
}

/// Read the latest value for a key hash at or before `max_version`.
pub(crate) fn read_value(
    txn: &redb::ReadTransaction,
    max_version: Version,
    key_hash: KeyHash,
) -> Result<Option<OwnedValue>> {
    let values = txn.open_table(VALUES_TABLE)?;
    let tombstones = txn.open_table(TOMBSTONES_TABLE)?;
    resolve_value(&values, &tombstones, max_version, key_hash)
}

/// Read the current unified JMT version from metadata, if any.
pub(crate) fn current_version(txn: &redb::ReadTransaction) -> Result<Option<Version>> {
    let table = txn.open_table(METADATA_TABLE)?;
    match table.get(JMT_VERSION_KEY)? {
        Some(guard) => {
            let bytes: [u8; 8] = guard
                .value()
                .try_into()
                .map_err(|_| StateTreeError::Serialization("invalid version bytes in metadata".into()))?;
            Ok(Some(u64::from_le_bytes(bytes)))
        }
        None => Ok(None),
    }
}

fn version_from_key(vk: &[u8]) -> Result<Version> {
    if vk.len() < 40 {
        return Err(StateTreeError::Serialization(format!(
            "version key must be at least 40 bytes, got {}",
            vk.len()
        )));
    }
    let mut buf = [0u8; 8];
    buf.copy_from_slice(&vk[32..40]);
    Ok(u64::from_be_bytes(buf))
}

fn version_key(key_hash: &KeyHash, version: Version) -> [u8; 40] {
    let mut vk = [0u8; 40];
    vk[..32].copy_from_slice(&key_hash.0);
    vk[32..].copy_from_slice(&version.to_be_bytes());
    vk
}

fn version_range_keys(key_hash: &KeyHash, max_version: Version) -> ([u8; 40], [u8; 40]) {
    let start = version_key(key_hash, 0);
    let end = version_key(key_hash, max_version);
    (start, end)
}

#[cfg(test)]
mod tests {
    use super::*;

    fn test_store() -> RedbNodeStore {
        let dir = tempfile::tempdir().unwrap();
        RedbNodeStore::open(dir.path().join("test.redb")).unwrap()
    }

    #[test]
    fn open_creates_tables() {
        let store = test_store();
        store
            .read(|txn| {
                txn.open_table(NODES_TABLE)?;
                txn.open_table(VALUES_TABLE)?;
                txn.open_table(TOMBSTONES_TABLE)?;
                txn.open_table(METADATA_TABLE)?;
                Ok(())
            })
            .unwrap();
    }

    #[test]
    fn large_value_roundtrip() {
        let store = test_store();
        let key_hash = KeyHash([0xab; 32]);
        let payload = vec![0x42u8; 2 * 1024 * 1024];

        let wtxn = store.begin_write().unwrap();
        write_value(&wtxn, 0, key_hash, Some(payload.clone())).unwrap();
        wtxn.commit().unwrap();

        store
            .read(|txn| {
                let val = read_value(txn, 0, key_hash)?.expect("value should exist");
                assert_eq!(val.len(), payload.len());
                assert_eq!(val, payload);
                Ok(())
            })
            .unwrap();
    }

    /// Regression: delete the rightmost leaf then insert a smaller leaf. The
    /// metadata side-table (`existing_was_invalidated` path in `write_nodes`)
    /// must not leave `get_rightmost_leaf` pointing at the deleted node.
    #[test]
    fn rightmost_leaf_cleared_on_delete_then_smaller_insert() {
        use jmt::{storage::TreeReader, JellyfishMerkleTree};

        let store = test_store();

        let kh_a = KeyHash([0xff; 32]); // larger
        let kh_b = KeyHash([0x11; 32]); // smaller

        // Insert leaf A — becomes the rightmost.
        let rtxn = store.begin_read().unwrap();
        let reader = TransactionReader::new(rtxn).unwrap();
        let tree = JellyfishMerkleTree::<_, sha2::Sha256>::new(&reader);
        let (_, batch) = tree.put_value_set(vec![(kh_a, Some(vec![1]))], 0).unwrap();
        drop(reader);

        let wtxn = store.begin_write().unwrap();
        write_nodes(&wtxn, &batch).unwrap();
        wtxn.commit().unwrap();

        let rtxn = store.begin_read().unwrap();
        let reader = TransactionReader::new(rtxn).unwrap();
        let (_, leaf) = reader.get_rightmost_leaf().unwrap().expect("should exist after insert");
        assert_eq!(leaf.key_hash(), kh_a);

        // Delete A (set to None) and insert B with a smaller key_hash — single batch.
        let tree = JellyfishMerkleTree::<_, sha2::Sha256>::new(&reader);
        let (_, batch) = tree
            .put_value_set(vec![(kh_a, None), (kh_b, Some(vec![2]))], 1)
            .unwrap();
        drop(reader);

        let wtxn = store.begin_write().unwrap();
        write_nodes(&wtxn, &batch).unwrap();
        wtxn.commit().unwrap();

        // Assert: rightmost must be B or None — never the deleted A.
        let rtxn = store.begin_read().unwrap();
        let reader = TransactionReader::new(rtxn).unwrap();
        let rightmost = reader.get_rightmost_leaf().unwrap();
        if let Some((_, leaf)) = rightmost {
            assert_ne!(leaf.key_hash(), kh_a, "rightmost must not point at deleted A");
            assert_eq!(leaf.key_hash(), kh_b, "rightmost must be the live leaf B");
        }
    }

    #[test]
    fn tombstone_then_reinsert() {
        let store = test_store();
        let key_hash = KeyHash([0xab; 32]);

        let wtxn = store.begin_write().unwrap();
        write_value(&wtxn, 0, key_hash, Some(vec![1])).unwrap();
        write_value(&wtxn, 1, key_hash, None).unwrap();
        write_value(&wtxn, 2, key_hash, Some(vec![2])).unwrap();
        wtxn.commit().unwrap();

        store
            .read(|txn| {
                assert_eq!(read_value(txn, 0, key_hash)?, Some(vec![1]));
                assert_eq!(read_value(txn, 1, key_hash)?, None);
                assert_eq!(read_value(txn, 2, key_hash)?, Some(vec![2]));
                Ok(())
            })
            .unwrap();
    }
}
