use crate::error::AedbError;
use crate::storage::keyspace::KvEntry;
use parking_lot::Mutex;
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use std::collections::{BTreeMap, HashMap, HashSet};
use std::fs;
use std::fs::File;
use std::io::{Read, Seek, SeekFrom, Write};
use std::ops::Bound;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{SystemTime, UNIX_EPOCH};
use tempfile::NamedTempFile;

const KV_SEGMENT_DIR: &str = "kv_segments";
const KV_SEGMENT_MAGIC: &[u8; 8] = b"AEDBKV1\0";
const KV_SEGMENT_BLOCK_ENTRIES: usize = 64;
const DEFAULT_KV_SEGMENT_BLOCK_CACHE_BYTES: usize = 16 * 1024 * 1024;
const KV_SEGMENT_BLOOM_WORDS: usize = 32;
const KV_SEGMENT_BLOOM_BITS: u64 = (KV_SEGMENT_BLOOM_WORDS as u64) * 64;
const KV_SEGMENT_BLOOM_HASHES: u64 = 4;

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct KvSegmentMeta {
    pub filename: String,
    pub min_key: Vec<u8>,
    pub max_key: Vec<u8>,
    pub entry_count: u64,
    pub sha256_hex: String,
    pub created_at_micros: u64,
    #[serde(default)]
    pub bloom_bits: Vec<u64>,
    #[serde(default)]
    pub blocks: Vec<KvSegmentBlockMeta>,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct KvSegmentBlockMeta {
    pub first_key: Vec<u8>,
    pub last_key: Vec<u8>,
    pub byte_start: u64,
    pub byte_count: u64,
    pub entry_count: u32,
    pub sha256_hex: String,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct KvSegmentEntry {
    pub key: Vec<u8>,
    pub entry: KvEntry,
}

#[derive(Debug)]
pub struct KvSegmentStore {
    dir: PathBuf,
    block_cache_capacity_bytes: usize,
    block_cache: Mutex<KvSegmentBlockCache>,
    /// Cumulative block-cache lookups served from memory.
    block_cache_hits: AtomicU64,
    /// Cumulative block-cache lookups that fell through to a segment read.
    block_cache_misses: AtomicU64,
    pending_publish_filenames: Mutex<HashSet<String>>,
}

impl PartialEq for KvSegmentStore {
    fn eq(&self, other: &Self) -> bool {
        self.dir == other.dir
    }
}

impl Eq for KvSegmentStore {}

impl KvSegmentStore {
    pub fn open(data_dir: &Path) -> Result<Self, AedbError> {
        Self::open_with_block_cache_bytes(data_dir, DEFAULT_KV_SEGMENT_BLOCK_CACHE_BYTES)
    }

    pub fn open_with_block_cache_bytes(
        data_dir: &Path,
        block_cache_capacity_bytes: usize,
    ) -> Result<Self, AedbError> {
        let dir = data_dir.join(KV_SEGMENT_DIR);
        fs::create_dir_all(&dir)?;
        Ok(Self {
            dir,
            block_cache_capacity_bytes,
            block_cache: Mutex::new(KvSegmentBlockCache::new(block_cache_capacity_bytes)),
            block_cache_hits: AtomicU64::new(0),
            block_cache_misses: AtomicU64::new(0),
            pending_publish_filenames: Mutex::new(HashSet::new()),
        })
    }

    pub fn block_cache_resident_bytes(&self) -> usize {
        self.block_cache.lock().resident_bytes()
    }

    /// Cumulative block-cache hits since the store was opened.
    pub fn block_cache_hits(&self) -> u64 {
        self.block_cache_hits.load(Ordering::Relaxed)
    }

    /// Cumulative block-cache misses since the store was opened.
    pub fn block_cache_misses(&self) -> u64 {
        self.block_cache_misses.load(Ordering::Relaxed)
    }

    pub fn block_cache_capacity_bytes(&self) -> usize {
        self.block_cache_capacity_bytes
    }

    pub fn list_segment_filenames(&self) -> Result<Vec<String>, AedbError> {
        let mut filenames = Vec::new();
        for entry in fs::read_dir(&self.dir)? {
            let entry = entry?;
            let filename = entry.file_name().to_string_lossy().to_string();
            if !is_kv_segment_filename(&filename) {
                continue;
            }
            filenames.push(filename);
        }
        filenames.sort();
        Ok(filenames)
    }

    pub fn reclaim_unreferenced_segments(
        &self,
        referenced_filenames: &HashSet<String>,
    ) -> Result<usize, AedbError> {
        let mut reclaimed_count = 0usize;
        let pending_publish_filenames = self.pending_publish_filenames.lock();
        for filename in self.list_segment_filenames()? {
            if referenced_filenames.contains(&filename)
                || pending_publish_filenames.contains(&filename)
            {
                continue;
            }
            validate_segment_filename(&filename)?;
            self.block_cache.lock().remove_filename(&filename);
            fs::remove_file(self.dir.join(&filename))?;
            reclaimed_count = reclaimed_count.saturating_add(1);
        }
        Ok(reclaimed_count)
    }

    pub fn write_segment(
        &self,
        namespace_hint: &str,
        entries: Vec<KvSegmentEntry>,
    ) -> Result<KvSegmentMeta, AedbError> {
        if entries.is_empty() {
            return Err(AedbError::Validation(
                "cannot write empty KV segment".into(),
            ));
        }
        for pair in entries.windows(2) {
            if pair[0].key >= pair[1].key {
                return Err(AedbError::Validation(
                    "KV segment entries must be strictly sorted".into(),
                ));
            }
        }

        let min_key = entries[0].key.clone();
        let max_key = entries[entries.len() - 1].key.clone();
        let bloom_bits = build_bloom(&entries);
        let mut payload = Vec::with_capacity(
            KV_SEGMENT_MAGIC.len()
                + entries
                    .iter()
                    .map(|entry| entry.key.len().saturating_add(entry.entry.value.len() + 64))
                    .sum::<usize>(),
        );
        payload.extend_from_slice(KV_SEGMENT_MAGIC);
        let mut blocks = Vec::with_capacity(entries.len().div_ceil(KV_SEGMENT_BLOCK_ENTRIES));
        for block in entries.chunks(KV_SEGMENT_BLOCK_ENTRIES) {
            let byte_start = payload.len() as u64;
            for item in block {
                let encoded_entry =
                    rmp_serde::to_vec(&item.entry).map_err(|e| AedbError::Encode(e.to_string()))?;
                write_u32(&mut payload, item.key.len(), "KV segment key")?;
                write_u32(&mut payload, encoded_entry.len(), "KV segment entry")?;
                payload.extend_from_slice(&item.key);
                payload.extend_from_slice(&encoded_entry);
            }
            let byte_count = (payload.len() as u64).saturating_sub(byte_start);
            let block_start = byte_start as usize;
            let block_end = payload.len();
            blocks.push(KvSegmentBlockMeta {
                first_key: block[0].key.clone(),
                last_key: block[block.len() - 1].key.clone(),
                byte_start,
                byte_count,
                entry_count: block.len() as u32,
                sha256_hex: hex_string(Sha256::digest(&payload[block_start..block_end]).as_slice()),
            });
        }
        let sha256_hex = hex_string(Sha256::digest(&payload).as_slice());
        let created_at_micros = now_micros();
        let filename = format!(
            "kvseg_{}_{}_{}.aedbkv",
            sanitize_hint(namespace_hint),
            created_at_micros,
            uuid::Uuid::new_v4().simple()
        );
        let final_path = self.dir.join(&filename);
        let mut tmp = NamedTempFile::new_in(&self.dir)?;
        tmp.write_all(&payload)?;
        tmp.flush()?;
        tmp.as_file().sync_all()?;
        self.pending_publish_filenames
            .lock()
            .insert(filename.clone());
        if let Err(e) = tmp.persist(&final_path) {
            self.pending_publish_filenames.lock().remove(&filename);
            return Err(AedbError::Io(e.error));
        }

        Ok(KvSegmentMeta {
            filename,
            min_key,
            max_key,
            entry_count: entries.len() as u64,
            sha256_hex,
            created_at_micros,
            bloom_bits,
            blocks,
        })
    }

    pub(crate) fn mark_segment_published(&self, filename: &str) {
        self.pending_publish_filenames.lock().remove(filename);
    }

    pub fn read_segment(&self, meta: &KvSegmentMeta) -> Result<Vec<KvSegmentEntry>, AedbError> {
        self.read_segment_with_cache_policy(meta, true)
    }

    pub fn read_segment_cold(
        &self,
        meta: &KvSegmentMeta,
    ) -> Result<Vec<KvSegmentEntry>, AedbError> {
        self.read_segment_with_cache_policy(meta, false)
    }

    fn read_segment_with_cache_policy(
        &self,
        meta: &KvSegmentMeta,
        populate_block_cache: bool,
    ) -> Result<Vec<KvSegmentEntry>, AedbError> {
        validate_segment_filename(&meta.filename)?;
        if meta.blocks.is_empty() {
            return self.read_legacy_segment(meta);
        }
        let mut file = self.open_validated_segment_file(&meta.filename)?;
        let mut entries = Vec::with_capacity(meta.entry_count as usize);
        for block in &meta.blocks {
            let cache_key = block_cache_key(meta, block);
            let block_entries = if populate_block_cache
                && let Some(cached) = self.block_cache.lock().get(&cache_key)
            {
                self.block_cache_hits.fetch_add(1, Ordering::Relaxed);
                cached
            } else {
                let decoded = read_block_entries_from_open_file(block, &mut file)?;
                if populate_block_cache {
                    self.block_cache_misses.fetch_add(1, Ordering::Relaxed);
                    self.block_cache
                        .lock()
                        .insert(cache_key, Arc::clone(&decoded));
                }
                decoded
            };
            entries.extend(block_entries.iter().cloned());
        }
        validate_segment_entries(meta, &entries)?;
        Ok(entries)
    }

    pub fn scan_range(
        &self,
        meta: &KvSegmentMeta,
        start: &Bound<Vec<u8>>,
        end: &Bound<Vec<u8>>,
    ) -> Result<Vec<KvSegmentEntry>, AedbError> {
        self.scan_range_with_cache_policy(meta, start, end, usize::MAX, true)
    }

    pub(crate) fn scan_range_limited(
        &self,
        meta: &KvSegmentMeta,
        start: &Bound<Vec<u8>>,
        end: &Bound<Vec<u8>>,
        limit: usize,
    ) -> Result<Vec<KvSegmentEntry>, AedbError> {
        self.scan_range_with_cache_policy(meta, start, end, limit, true)
    }

    pub(crate) fn scan_range_limited_cold(
        &self,
        meta: &KvSegmentMeta,
        start: &Bound<Vec<u8>>,
        end: &Bound<Vec<u8>>,
        limit: usize,
    ) -> Result<Vec<KvSegmentEntry>, AedbError> {
        self.scan_range_with_cache_policy(meta, start, end, limit, false)
    }

    pub fn scan_range_cold(
        &self,
        meta: &KvSegmentMeta,
        start: &Bound<Vec<u8>>,
        end: &Bound<Vec<u8>>,
    ) -> Result<Vec<KvSegmentEntry>, AedbError> {
        self.scan_range_with_cache_policy(meta, start, end, usize::MAX, false)
    }

    fn scan_range_with_cache_policy(
        &self,
        meta: &KvSegmentMeta,
        start: &Bound<Vec<u8>>,
        end: &Bound<Vec<u8>>,
        limit: usize,
        populate_block_cache: bool,
    ) -> Result<Vec<KvSegmentEntry>, AedbError> {
        if limit == 0 {
            return Ok(Vec::new());
        }
        if !segment_overlaps_bounds(meta, start, end) {
            return Ok(Vec::new());
        }
        if meta.blocks.is_empty() {
            return self
                .read_legacy_segment(meta)
                .map(|entries| filter_entries_by_bounds(entries, start, end, limit));
        }
        let mut entries = Vec::with_capacity(limit.min(KV_SEGMENT_BLOCK_ENTRIES));
        let first_block_position = first_block_position_for_start(&meta.blocks, start);
        let mut file = None;
        for block in &meta.blocks[first_block_position..] {
            if block_starts_after_end(block, end) {
                break;
            }
            let block_entries = self.read_block_entries_with_shared_file(
                meta,
                block,
                populate_block_cache,
                &mut file,
            )?;
            for entry in block_entries
                .iter()
                .filter(|entry| key_within_bounds(&entry.key, start, end))
            {
                entries.push(entry.clone());
                if entries.len() >= limit {
                    return Ok(entries);
                }
            }
        }
        Ok(entries)
    }

    fn read_block_entries_with_shared_file(
        &self,
        meta: &KvSegmentMeta,
        block: &KvSegmentBlockMeta,
        populate_block_cache: bool,
        file: &mut Option<File>,
    ) -> Result<Arc<[KvSegmentEntry]>, AedbError> {
        let cache_key = block_cache_key(meta, block);
        if populate_block_cache && let Some(entries) = self.block_cache.lock().get(&cache_key) {
            return Ok(entries);
        }
        if file.is_none() {
            *file = Some(self.open_validated_segment_file(&meta.filename)?);
        }
        let decoded =
            read_block_entries_from_open_file(block, file.as_mut().expect("segment file opened"))?;
        if populate_block_cache {
            self.block_cache
                .lock()
                .insert(cache_key, Arc::clone(&decoded));
        }
        Ok(decoded)
    }

    fn read_legacy_segment(&self, meta: &KvSegmentMeta) -> Result<Vec<KvSegmentEntry>, AedbError> {
        let payload = fs::read(self.dir.join(&meta.filename))?;
        let hash = hex_string(Sha256::digest(&payload).as_slice());
        if hash != meta.sha256_hex {
            return Err(AedbError::IntegrityError {
                message: "KV segment checksum mismatch".into(),
            });
        }
        if payload.len() < KV_SEGMENT_MAGIC.len()
            || &payload[..KV_SEGMENT_MAGIC.len()] != KV_SEGMENT_MAGIC
        {
            return Err(AedbError::IntegrityError {
                message: "KV segment magic mismatch".into(),
            });
        }
        let entries: Vec<KvSegmentEntry> =
            rmp_serde::from_slice(&payload[KV_SEGMENT_MAGIC.len()..])
                .map_err(|e| AedbError::Decode(e.to_string()))?;
        validate_segment_entries(meta, &entries)?;
        Ok(entries)
    }

    fn read_block_entries(
        &self,
        meta: &KvSegmentMeta,
        block: &KvSegmentBlockMeta,
    ) -> Result<Arc<[KvSegmentEntry]>, AedbError> {
        self.read_block_entries_with_cache_policy(meta, block, true)
    }

    fn read_block_entries_with_cache_policy(
        &self,
        meta: &KvSegmentMeta,
        block: &KvSegmentBlockMeta,
        populate_block_cache: bool,
    ) -> Result<Arc<[KvSegmentEntry]>, AedbError> {
        validate_segment_filename(&meta.filename)?;
        let cache_key = block_cache_key(meta, block);
        if populate_block_cache && let Some(entries) = self.block_cache.lock().get(&cache_key) {
            return Ok(entries);
        }
        let mut file = self.open_validated_segment_file(&meta.filename)?;
        let entries = read_block_entries_from_open_file(block, &mut file)?;
        if populate_block_cache {
            self.block_cache
                .lock()
                .insert(cache_key, Arc::clone(&entries));
        }
        Ok(entries)
    }

    fn open_validated_segment_file(&self, filename: &str) -> Result<File, AedbError> {
        let mut file = File::open(self.dir.join(filename))?;
        let mut magic = [0u8; KV_SEGMENT_MAGIC.len()];
        file.read_exact(&mut magic)?;
        if &magic != KV_SEGMENT_MAGIC {
            return Err(AedbError::IntegrityError {
                message: "KV segment magic mismatch".into(),
            });
        }
        Ok(file)
    }

    pub fn get(&self, meta: &KvSegmentMeta, key: &[u8]) -> Result<Option<KvEntry>, AedbError> {
        if key < meta.min_key.as_slice() || key > meta.max_key.as_slice() {
            return Ok(None);
        }
        if !bloom_may_contain(meta, key) {
            return Ok(None);
        }
        if meta.blocks.is_empty() {
            let entries = self.read_legacy_segment(meta)?;
            return Ok(entries
                .binary_search_by(|entry| entry.key.as_slice().cmp(key))
                .ok()
                .map(|entry_position| entries[entry_position].entry.clone()));
        }
        let block_position = match meta
            .blocks
            .binary_search_by(|block| block.last_key.as_slice().cmp(key))
        {
            Ok(found_position) => found_position,
            Err(candidate_position) => candidate_position,
        };
        let Some(block) = meta.blocks.get(block_position) else {
            return Ok(None);
        };
        if key < block.first_key.as_slice() || key > block.last_key.as_slice() {
            return Ok(None);
        }
        let entries = self.read_block_entries(meta, block)?;
        Ok(entries
            .as_ref()
            .binary_search_by(|entry| entry.key.as_slice().cmp(key))
            .ok()
            .map(|entry_position| entries[entry_position].entry.clone()))
    }
}

fn block_cache_key(meta: &KvSegmentMeta, block: &KvSegmentBlockMeta) -> KvSegmentBlockCacheKey {
    KvSegmentBlockCacheKey {
        filename: meta.filename.clone(),
        byte_start: block.byte_start,
        byte_count: block.byte_count,
        sha256_hex: block.sha256_hex.clone(),
    }
}

fn read_block_entries_from_open_file(
    block: &KvSegmentBlockMeta,
    file: &mut File,
) -> Result<Arc<[KvSegmentEntry]>, AedbError> {
    file.seek(SeekFrom::Start(block.byte_start))?;
    let block_bytes = usize::try_from(block.byte_count).map_err(|_| {
        AedbError::Validation("KV segment block is too large to read on this platform".into())
    })?;
    let mut payload = vec![0u8; block_bytes];
    file.read_exact(&mut payload)?;
    let hash = hex_string(Sha256::digest(&payload).as_slice());
    if hash != block.sha256_hex {
        return Err(AedbError::IntegrityError {
            message: "KV segment block checksum mismatch".into(),
        });
    }
    let mut cursor = std::io::Cursor::new(payload);
    let mut entries = Vec::with_capacity(block.entry_count as usize);
    while cursor.position() < block.byte_count {
        let key_bytes = read_u32(&mut cursor)? as usize;
        let entry_bytes = read_u32(&mut cursor)? as usize;
        let mut key = vec![0u8; key_bytes];
        cursor.read_exact(&mut key)?;
        let mut encoded_entry = vec![0u8; entry_bytes];
        cursor.read_exact(&mut encoded_entry)?;
        let entry: KvEntry =
            rmp_serde::from_slice(&encoded_entry).map_err(|e| AedbError::Decode(e.to_string()))?;
        entries.push(KvSegmentEntry { key, entry });
    }
    if entries.len() as u32 != block.entry_count {
        return Err(AedbError::IntegrityError {
            message: "KV segment block entry count mismatch".into(),
        });
    }
    if entries.first().map(|entry| &entry.key) != Some(&block.first_key)
        || entries.last().map(|entry| &entry.key) != Some(&block.last_key)
    {
        return Err(AedbError::IntegrityError {
            message: "KV segment block key bounds mismatch".into(),
        });
    }
    for pair in entries.windows(2) {
        if pair[0].key >= pair[1].key {
            return Err(AedbError::IntegrityError {
                message: "KV segment block keys are not strictly sorted".into(),
            });
        }
    }
    let entries: Arc<[KvSegmentEntry]> = Arc::from(entries);
    Ok(entries)
}

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
struct KvSegmentBlockCacheKey {
    filename: String,
    byte_start: u64,
    byte_count: u64,
    sha256_hex: String,
}

#[derive(Debug)]
struct CachedKvSegmentBlock {
    entries: Arc<[KvSegmentEntry]>,
    resident_bytes: usize,
    access_tick: u64,
}

#[derive(Debug)]
struct KvSegmentBlockCache {
    capacity_bytes: usize,
    resident_bytes: usize,
    access_tick: u64,
    blocks: HashMap<KvSegmentBlockCacheKey, CachedKvSegmentBlock>,
    access_order: BTreeMap<u64, KvSegmentBlockCacheKey>,
}

impl KvSegmentBlockCache {
    fn new(capacity_bytes: usize) -> Self {
        Self {
            capacity_bytes,
            resident_bytes: 0,
            access_tick: 0,
            blocks: HashMap::new(),
            access_order: BTreeMap::new(),
        }
    }

    fn resident_bytes(&self) -> usize {
        self.resident_bytes
    }

    fn next_access_tick(&mut self) -> u64 {
        self.access_tick = self.access_tick.saturating_add(1);
        self.access_tick
    }

    fn get(&mut self, cache_key: &KvSegmentBlockCacheKey) -> Option<Arc<[KvSegmentEntry]>> {
        let access_tick = self.next_access_tick();
        let cached = self.blocks.get_mut(cache_key)?;
        self.access_order.remove(&cached.access_tick);
        cached.access_tick = access_tick;
        self.access_order.insert(access_tick, cache_key.clone());
        Some(Arc::clone(&cached.entries))
    }

    fn insert(&mut self, cache_key: KvSegmentBlockCacheKey, entries: Arc<[KvSegmentEntry]>) {
        let resident_bytes = kv_segment_block_entries_cost(&entries);
        if self.capacity_bytes == 0 || resident_bytes > self.capacity_bytes {
            return;
        }
        if let Some(cached) = self.blocks.remove(&cache_key) {
            self.access_order.remove(&cached.access_tick);
            self.resident_bytes = self.resident_bytes.saturating_sub(cached.resident_bytes);
        }
        let access_tick = self.next_access_tick();
        self.resident_bytes = self.resident_bytes.saturating_add(resident_bytes);
        self.access_order.insert(access_tick, cache_key.clone());
        self.blocks.insert(
            cache_key,
            CachedKvSegmentBlock {
                entries,
                resident_bytes,
                access_tick,
            },
        );
        self.evict_to_capacity();
    }

    fn remove_filename(&mut self, filename: &str) {
        let keys = self
            .blocks
            .keys()
            .filter(|cache_key| cache_key.filename == filename)
            .cloned()
            .collect::<Vec<_>>();
        for cache_key in keys {
            if let Some(cached) = self.blocks.remove(&cache_key) {
                self.access_order.remove(&cached.access_tick);
                self.resident_bytes = self.resident_bytes.saturating_sub(cached.resident_bytes);
            }
        }
    }

    fn evict_to_capacity(&mut self) {
        while self.resident_bytes > self.capacity_bytes {
            let Some((access_tick, cache_key)) = self
                .access_order
                .iter()
                .next()
                .map(|(access_tick, cache_key)| (*access_tick, cache_key.clone()))
            else {
                break;
            };
            self.access_order.remove(&access_tick);
            if let Some(cached) = self.blocks.remove(&cache_key) {
                self.resident_bytes = self.resident_bytes.saturating_sub(cached.resident_bytes);
            }
        }
    }
}

fn kv_segment_block_entries_cost(entries: &[KvSegmentEntry]) -> usize {
    entries
        .iter()
        .map(|entry| {
            entry
                .key
                .len()
                .saturating_add(entry.entry.value.len())
                .saturating_add(96)
        })
        .sum()
}

fn validate_segment_entries(
    meta: &KvSegmentMeta,
    entries: &[KvSegmentEntry],
) -> Result<(), AedbError> {
    if entries.len() as u64 != meta.entry_count {
        return Err(AedbError::IntegrityError {
            message: "KV segment entry count mismatch".into(),
        });
    }
    if entries.first().map(|entry| &entry.key) != Some(&meta.min_key)
        || entries.last().map(|entry| &entry.key) != Some(&meta.max_key)
    {
        return Err(AedbError::IntegrityError {
            message: "KV segment key bounds mismatch".into(),
        });
    }
    for pair in entries.windows(2) {
        if pair[0].key >= pair[1].key {
            return Err(AedbError::IntegrityError {
                message: "KV segment entries are not strictly sorted".into(),
            });
        }
    }
    Ok(())
}

fn filter_entries_by_bounds(
    entries: Vec<KvSegmentEntry>,
    start: &Bound<Vec<u8>>,
    end: &Bound<Vec<u8>>,
    limit: usize,
) -> Vec<KvSegmentEntry> {
    entries
        .into_iter()
        .filter(|entry| key_within_bounds(&entry.key, start, end))
        .take(limit)
        .collect()
}

fn key_within_bounds(key: &[u8], start: &Bound<Vec<u8>>, end: &Bound<Vec<u8>>) -> bool {
    let after_start = match start {
        Bound::Included(start) => key >= start.as_slice(),
        Bound::Excluded(start) => key > start.as_slice(),
        Bound::Unbounded => true,
    };
    let before_end = match end {
        Bound::Included(end) => key <= end.as_slice(),
        Bound::Excluded(end) => key < end.as_slice(),
        Bound::Unbounded => true,
    };
    after_start && before_end
}

fn segment_overlaps_bounds(
    meta: &KvSegmentMeta,
    start: &Bound<Vec<u8>>,
    end: &Bound<Vec<u8>>,
) -> bool {
    let after_start = match start {
        Bound::Included(start) => meta.max_key.as_slice() >= start.as_slice(),
        Bound::Excluded(start) => meta.max_key.as_slice() > start.as_slice(),
        Bound::Unbounded => true,
    };
    let before_end = match end {
        Bound::Included(end) => meta.min_key.as_slice() <= end.as_slice(),
        Bound::Excluded(end) => meta.min_key.as_slice() < end.as_slice(),
        Bound::Unbounded => true,
    };
    after_start && before_end
}

fn first_block_position_for_start(blocks: &[KvSegmentBlockMeta], start: &Bound<Vec<u8>>) -> usize {
    match start {
        Bound::Unbounded => 0,
        Bound::Included(start) => {
            match blocks.binary_search_by(|block| block.last_key.as_slice().cmp(start.as_slice())) {
                Ok(position) | Err(position) => position,
            }
        }
        Bound::Excluded(start) => {
            match blocks.binary_search_by(|block| block.last_key.as_slice().cmp(start.as_slice())) {
                Ok(position) => position.saturating_add(1),
                Err(position) => position,
            }
        }
    }
}

fn block_starts_after_end(block: &KvSegmentBlockMeta, end: &Bound<Vec<u8>>) -> bool {
    match end {
        Bound::Unbounded => false,
        Bound::Included(end) => block.first_key.as_slice() > end.as_slice(),
        Bound::Excluded(end) => block.first_key.as_slice() >= end.as_slice(),
    }
}

fn write_u32(out: &mut Vec<u8>, value: usize, label: &str) -> Result<(), AedbError> {
    let value = u32::try_from(value)
        .map_err(|_| AedbError::Validation(format!("{label} exceeds u32 frame bound")))?;
    out.extend_from_slice(&value.to_le_bytes());
    Ok(())
}

fn read_u32(cursor: &mut std::io::Cursor<Vec<u8>>) -> Result<u32, AedbError> {
    let mut bytes = [0u8; 4];
    cursor.read_exact(&mut bytes)?;
    Ok(u32::from_le_bytes(bytes))
}

fn build_bloom(entries: &[KvSegmentEntry]) -> Vec<u64> {
    let mut bits = vec![0u64; KV_SEGMENT_BLOOM_WORDS];
    for entry in entries {
        for bit in bloom_bits_for_key(&entry.key) {
            let word = (bit / 64) as usize;
            let bit_shift = bit % 64;
            bits[word] |= 1u64 << bit_shift;
        }
    }
    bits
}

fn bloom_may_contain(meta: &KvSegmentMeta, key: &[u8]) -> bool {
    if meta.bloom_bits.len() != KV_SEGMENT_BLOOM_WORDS {
        return true;
    }
    bloom_bits_for_key(key).into_iter().all(|bit| {
        let word = (bit / 64) as usize;
        let bit_shift = bit % 64;
        meta.bloom_bits[word] & (1u64 << bit_shift) != 0
    })
}

fn bloom_bits_for_key(key: &[u8]) -> [u64; KV_SEGMENT_BLOOM_HASHES as usize] {
    let digest = blake3::hash(key);
    let bytes = digest.as_bytes();
    let first = u64::from_le_bytes(copy_hash_bytes::<8>(bytes, 0));
    let second = u64::from_le_bytes(copy_hash_bytes::<8>(bytes, 8)) | 1;
    let mut out = [0u64; KV_SEGMENT_BLOOM_HASHES as usize];
    for probe_number in 0..KV_SEGMENT_BLOOM_HASHES {
        out[probe_number as usize] =
            first.wrapping_add(probe_number.wrapping_mul(second)) % KV_SEGMENT_BLOOM_BITS;
    }
    out
}

fn copy_hash_bytes<const N: usize>(bytes: &[u8; 32], offset: usize) -> [u8; N] {
    let mut out = [0u8; N];
    out.copy_from_slice(&bytes[offset..offset + N]);
    out
}

fn validate_segment_filename(filename: &str) -> Result<(), AedbError> {
    if filename.is_empty() || filename.contains('/') || filename.contains('\\') {
        return Err(AedbError::Validation("invalid KV segment filename".into()));
    }
    Ok(())
}

fn is_kv_segment_filename(filename: &str) -> bool {
    filename.starts_with("kvseg_")
        && filename.ends_with(".aedbkv")
        && validate_segment_filename(filename).is_ok()
}

fn sanitize_hint(input: &str) -> String {
    let mut out = String::with_capacity(input.len().min(48));
    for b in input.bytes().take(48) {
        if b.is_ascii_alphanumeric() {
            out.push(b as char);
        } else {
            out.push('_');
        }
    }
    if out.is_empty() { "ns".into() } else { out }
}

fn now_micros() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_micros() as u64
}

fn hex_string(bytes: &[u8]) -> String {
    const HEX: &[u8; 16] = b"0123456789abcdef";
    let mut out = String::with_capacity(bytes.len() * 2);
    for byte in bytes {
        out.push(HEX[(byte >> 4) as usize] as char);
        out.push(HEX[(byte & 0x0f) as usize] as char);
    }
    out
}

#[cfg(test)]
mod tests;
