use crate::config::encode_namespaced_key;
use flaredb_storage::rocks_engine::RocksEngine;
use sha2::{Digest, Sha256};
use std::sync::Arc;

#[derive(Debug, Clone)]
pub struct MerkleLeaf {
    pub start: Vec<u8>,
    pub end: Vec<u8>,
    pub hash: Vec<u8>,
}

fn sha256(data: &[u8]) -> Vec<u8> {
    let mut hasher = Sha256::new();
    hasher.update(data);
    hasher.finalize().to_vec()
}

/// Build merkle tree for a namespace; returns (root, leaves).
pub fn build_merkle(
    engine: Arc<RocksEngine>,
    namespace_id: u32,
    chunk_size: usize,
) -> Result<(Vec<u8>, Vec<MerkleLeaf>), flaredb_storage::StorageError> {
    let prefix = namespace_id.to_be_bytes();
    let db = engine.db();
    let mut iter = db.iterator(rocksdb::IteratorMode::From(
        prefix.as_slice(),
        rocksdb::Direction::Forward,
    ));
    let mut leaves = Vec::new();
    let mut current_keys = Vec::new();
    let mut start_key = None;

    while let Some(Ok((k, v))) = iter.next() {
        if (k.len() < 4 || k[..4] != prefix) && !k.starts_with(&prefix) {
            break;
        }
        if start_key.is_none() {
            start_key = Some(k.to_vec());
        }
        // hash ts-decoded value to keep user value stable
        current_keys.push(sha256(&v));
        if current_keys.len() >= chunk_size {
            let leaf_hash = sha256(&current_keys.concat());
            let end = k.to_vec();
            leaves.push(MerkleLeaf {
                start: start_key.take().unwrap(),
                end: end.clone(),
                hash: leaf_hash,
            });
            current_keys.clear();
        }
    }

    if !current_keys.is_empty() {
        let leaf_hash = sha256(&current_keys.concat());
        let end = encode_namespaced_key(namespace_id, b"~end");
        leaves.push(MerkleLeaf {
            start: start_key.unwrap_or_else(|| encode_namespaced_key(namespace_id, b"")),
            end,
            hash: leaf_hash,
        });
    }

    if leaves.is_empty() {
        return Ok((sha256(b""), leaves));
    }

    let mut level: Vec<Vec<u8>> = leaves.iter().map(|l| l.hash.clone()).collect();
    while level.len() > 1 {
        let mut next = Vec::new();
        for chunk in level.chunks(2) {
            if chunk.len() == 2 {
                let mut c = Vec::new();
                c.extend_from_slice(&chunk[0]);
                c.extend_from_slice(&chunk[1]);
                next.push(sha256(&c));
            } else {
                next.push(chunk[0].clone());
            }
        }
        level = next;
    }

    Ok((level[0].clone(), leaves))
}

/// Fetch key/value pairs for a namespaced range.
pub async fn fetch_range(
    engine: Arc<RocksEngine>,
    namespace_id: u32,
    start: &[u8],
    end: &[u8],
) -> Result<(Vec<Vec<u8>>, Vec<Vec<u8>>), flaredb_storage::StorageError> {
    let mut keys = Vec::new();
    let mut values = Vec::new();
    let db = engine.db();
    let mut iter = db.iterator(rocksdb::IteratorMode::From(
        start,
        rocksdb::Direction::Forward,
    ));
    while let Some(Ok((k, v))) = iter.next() {
        if !k.starts_with(&namespace_id.to_be_bytes()) || (!end.is_empty() && k.as_ref() >= end) {
            break;
        }
        keys.push(k.to_vec());
        values.push(v.to_vec());
    }
    Ok((keys, values))
}
