//! DYF2 file format with footer-based index for efficient append operations.
//!
//! Layout:
//! ```text
//! ┌──────────────────────────────────┐
//! │ 8-byte Header: "DYF2" + flags   │
//! ├──────────────────────────────────┤
//! │ Arrow batches (base leaves)     │
//! │ Arrow batches (overflow/new)    │  ← append here
//! │ Arrow batches (field layers)    │  ← append here
//! ├──────────────────────────────────┤
//! │ FlatBuffers Index               │  ← rewrite only this + footer
//! ├──────────────────────────────────┤
//! │ 16-byte Footer                  │
//! │  - FB offset (u64 LE)           │
//! │  - FB size (u64 LE)             │
//! └──────────────────────────────────┘
//! ```
//!
//! Append cost: O(new_data + tree_metadata) instead of O(all_data).

#[allow(unused_imports, dead_code)]
mod schema {
    include!("schema/dyf_index_v2_generated.rs");
}

use schema::dyf::schema::v_2 as fb;

use std::collections::HashMap;
use std::fs::{File, OpenOptions};
use std::io::{self, Read, Seek, SeekFrom, Write};
use std::path::{Path, PathBuf};

use crate::tree::{LshTree, TreeNode};

/// Magic bytes for DYF2 format.
const DYF2_MAGIC: &[u8; 4] = b"DYF2";
/// Header size: 4 bytes magic + 4 bytes flags.
const HEADER_SIZE: u64 = 8;
/// Footer size: 8 bytes FB offset + 8 bytes FB size.
const FOOTER_SIZE: u64 = 16;

/// Descriptor for an Arrow IPC batch stored in the file.
#[derive(Clone, Debug)]
pub struct BatchDescriptor {
    /// Byte offset from data section start.
    pub offset: u64,
    /// Byte length of the Arrow IPC batch.
    pub length: u64,
    /// Number of rows in this batch.
    pub num_rows: u32,
}

/// A named field layer stored separately from base batches.
#[derive(Clone, Debug)]
pub struct FieldLayer {
    pub name: String,
    pub fields: Vec<FieldDescriptor>,
    pub batches: Vec<BatchDescriptor>,
}

/// Descriptor for a field within a layer.
#[derive(Clone, Debug)]
pub struct FieldDescriptor {
    pub name: String,
    pub arrow_type: String,
}

/// Build parameters stored in the index.
#[derive(Clone, Debug)]
pub struct BuildParams {
    pub max_depth: u8,
    pub num_bits: u8,
    pub min_leaf_size: u32,
    pub seed: u64,
    pub quantization: String,
    pub compression: String,
}

/// A DYF2 file handle supporting read, append, and rewrite operations.
pub struct DyfFile {
    path: PathBuf,
    tree: LshTree,
    base_batches: Vec<BatchDescriptor>,
    overflow_batches: Vec<BatchDescriptor>,
    field_layers: Vec<FieldLayer>,
    metadata: HashMap<String, String>,
    build_params: BuildParams,
    embedding_dim: u16,
    total_items: u64,
    /// Byte offset where new data can be appended (end of data section).
    data_end: u64,
}

impl DyfFile {
    /// Create a new DYF2 file.
    ///
    /// Writes the header, base Arrow batches (one per leaf), and the initial footer.
    /// `leaf_batches` contains the raw Arrow IPC bytes for each leaf, in order.
    pub fn create(
        path: &str,
        tree: LshTree,
        leaf_batches: &[Vec<u8>],
        build_params: BuildParams,
        metadata: HashMap<String, String>,
    ) -> io::Result<Self> {
        let embedding_dim = tree.embedding_dim() as u16;
        let total_items = tree.total_items() as u64;
        let leaf_indices = tree.leaf_indices();

        assert_eq!(
            leaf_batches.len(),
            leaf_indices.len(),
            "Must provide one batch per leaf"
        );

        let mut file = File::create(path)?;

        // Write header
        file.write_all(DYF2_MAGIC)?;
        file.write_all(&[0u8; 4])?; // flags (reserved)

        // Write base Arrow batches
        let mut base_batches = Vec::new();
        let mut offset = 0u64;

        for (i, batch_data) in leaf_batches.iter().enumerate() {
            let batch_desc = BatchDescriptor {
                offset,
                length: batch_data.len() as u64,
                num_rows: tree.node(leaf_indices[i]).num_items,
            };
            file.write_all(batch_data)?;
            offset += batch_data.len() as u64;
            base_batches.push(batch_desc);
        }

        let data_end = HEADER_SIZE + offset;

        let dyf = DyfFile {
            path: PathBuf::from(path),
            tree,
            base_batches,
            overflow_batches: Vec::new(),
            field_layers: Vec::new(),
            metadata,
            build_params,
            embedding_dim,
            total_items,
            data_end,
        };

        // Write footer (FlatBuffers index + 16-byte tail)
        dyf.write_footer(&mut file)?;

        Ok(dyf)
    }

    /// Open an existing DYF2 file.
    pub fn open(path: &str) -> io::Result<Self> {
        let mut file = File::open(path)?;
        let file_len = file.metadata()?.len();

        if file_len < HEADER_SIZE + FOOTER_SIZE {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                "File too small for DYF2",
            ));
        }

        // Read and verify magic
        let mut magic = [0u8; 4];
        file.read_exact(&mut magic)?;
        if &magic != DYF2_MAGIC {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                format!("Invalid magic: expected DYF2, got {:?}", &magic),
            ));
        }

        // Read footer (last 16 bytes)
        file.seek(SeekFrom::End(-16))?;
        let mut footer_buf = [0u8; 16];
        file.read_exact(&mut footer_buf)?;

        let fb_offset = u64::from_le_bytes(footer_buf[0..8].try_into().unwrap());
        let fb_size = u64::from_le_bytes(footer_buf[8..16].try_into().unwrap());

        // Read FlatBuffers index
        file.seek(SeekFrom::Start(fb_offset))?;
        let mut fb_buf = vec![0u8; fb_size as usize];
        file.read_exact(&mut fb_buf)?;

        let index = flatbuffers::root::<fb::Index>(&fb_buf).map_err(|e| {
            io::Error::new(
                io::ErrorKind::InvalidData,
                format!("FlatBuffers parse error: {e}"),
            )
        })?;

        // Parse tree nodes
        let embedding_dim = index.embedding_dim();
        let total_items = index.total_items();
        let num_bits = index
            .build_params()
            .map(|bp| bp.num_bits() as usize)
            .unwrap_or(3);
        let min_leaf_size = index
            .build_params()
            .map(|bp| bp.min_leaf_size() as usize)
            .unwrap_or(4);
        let seed = index.build_params().map(|bp| bp.seed()).unwrap_or(42);

        let tree = parse_tree_from_fb(&index, embedding_dim as usize, num_bits, min_leaf_size, seed);

        // Parse base batches
        let base_batches = parse_batch_descriptors(index.batches());

        // Parse overflow batches
        let overflow_batches = parse_batch_descriptors(index.overflow_batches());

        // Parse field layers
        let field_layers = parse_field_layers(index.field_layers());

        // Parse metadata
        let mut metadata = HashMap::new();
        if let Some(kvs) = index.metadata() {
            for i in 0..kvs.len() {
                let kv = kvs.get(i);
                if let (Some(k), Some(v)) = (kv.key(), kv.value()) {
                    metadata.insert(k.to_string(), v.to_string());
                }
            }
        }

        // Parse build params
        let build_params = if let Some(bp) = index.build_params() {
            BuildParams {
                max_depth: bp.max_depth(),
                num_bits: bp.num_bits(),
                min_leaf_size: bp.min_leaf_size(),
                seed: bp.seed(),
                quantization: bp.quantization().unwrap_or("float16").to_string(),
                compression: bp.compression().unwrap_or("none").to_string(),
            }
        } else {
            BuildParams {
                max_depth: 10,
                num_bits: num_bits as u8,
                min_leaf_size: min_leaf_size as u32,
                seed,
                quantization: "float16".into(),
                compression: "none".into(),
            }
        };

        let data_end = fb_offset;

        Ok(DyfFile {
            path: PathBuf::from(path),
            tree,
            base_batches,
            overflow_batches,
            field_layers,
            metadata,
            build_params,
            embedding_dim,
            total_items,
            data_end,
        })
    }

    /// Append a field layer to the file.
    ///
    /// `batch_data` contains Arrow IPC bytes for each batch in the layer.
    pub fn append_field_layer(
        &mut self,
        name: &str,
        fields: Vec<FieldDescriptor>,
        batch_data: &[Vec<u8>],
    ) -> io::Result<()> {
        let mut file = OpenOptions::new()
            .read(true)
            .write(true)
            .open(&self.path)?;

        // Seek to data_end (overwrite old footer)
        file.seek(SeekFrom::Start(self.data_end))?;

        // Write batch data
        let mut batches = Vec::new();
        let base_offset = self.data_end - HEADER_SIZE; // offset relative to data section
        let mut offset = base_offset;

        for batch in batch_data {
            batches.push(BatchDescriptor {
                offset,
                length: batch.len() as u64,
                num_rows: 0, // field layers don't track row counts independently
            });
            file.write_all(batch)?;
            offset += batch.len() as u64;
        }

        self.data_end = HEADER_SIZE + offset;

        self.field_layers.push(FieldLayer {
            name: name.to_string(),
            fields,
            batches,
        });

        // Rewrite footer
        self.write_footer(&mut file)?;

        Ok(())
    }

    /// Append overflow item batches to the file.
    pub fn append_items(&mut self, batch_data: &[Vec<u8>], num_rows: &[u32]) -> io::Result<()> {
        assert_eq!(batch_data.len(), num_rows.len());

        let mut file = OpenOptions::new()
            .read(true)
            .write(true)
            .open(&self.path)?;

        file.seek(SeekFrom::Start(self.data_end))?;

        let base_offset = self.data_end - HEADER_SIZE;
        let mut offset = base_offset;

        for (batch, &rows) in batch_data.iter().zip(num_rows.iter()) {
            self.overflow_batches.push(BatchDescriptor {
                offset,
                length: batch.len() as u64,
                num_rows: rows,
            });
            file.write_all(batch)?;
            offset += batch.len() as u64;
            self.total_items += rows as u64;
        }

        self.data_end = HEADER_SIZE + offset;

        self.write_footer(&mut file)?;

        Ok(())
    }

    /// Update metadata key-value pairs and rewrite the footer.
    pub fn update_metadata(&mut self, updates: &[(&str, &str)]) -> io::Result<()> {
        for &(k, v) in updates {
            self.metadata.insert(k.to_string(), v.to_string());
        }

        let mut file = OpenOptions::new()
            .read(true)
            .write(true)
            .open(&self.path)?;

        file.seek(SeekFrom::Start(self.data_end))?;
        self.write_footer(&mut file)?;

        Ok(())
    }

    /// Rewrite the footer after tree mutations.
    pub fn update_tree(&mut self) -> io::Result<()> {
        let mut file = OpenOptions::new()
            .read(true)
            .write(true)
            .open(&self.path)?;

        file.seek(SeekFrom::Start(self.data_end))?;
        self.write_footer(&mut file)?;

        Ok(())
    }

    /// Fragmentation ratio: overflow_items / total_items.
    pub fn fragmentation(&self) -> f32 {
        if self.total_items == 0 {
            return 0.0;
        }
        let overflow_items: u64 = self.overflow_batches.iter().map(|b| b.num_rows as u64).sum();
        overflow_items as f32 / self.total_items as f32
    }

    /// Compact the file: merge all overflow batches into base batches and
    /// rewrite the entire file. Field layers are preserved.
    ///
    /// This is a full rewrite — all data is read, concatenated, and written
    /// to a new file which replaces the original.
    pub fn compact(&mut self) -> io::Result<()> {
        if self.overflow_batches.is_empty() {
            return Ok(()); // nothing to compact
        }

        // Read all existing batch data
        let mut all_batch_data: Vec<Vec<u8>> = Vec::new();
        let mut all_num_rows: Vec<u32> = Vec::new();

        for desc in &self.base_batches {
            all_batch_data.push(self.read_batch_data(desc)?);
            all_num_rows.push(desc.num_rows);
        }
        for desc in &self.overflow_batches {
            all_batch_data.push(self.read_batch_data(desc)?);
            all_num_rows.push(desc.num_rows);
        }

        // Read field layer data
        let mut layer_data: Vec<Vec<Vec<u8>>> = Vec::new();
        for layer in &self.field_layers {
            let mut batches = Vec::new();
            for desc in &layer.batches {
                batches.push(self.read_batch_data(desc)?);
            }
            layer_data.push(batches);
        }

        // Write new file
        let tmp_path = self.path.with_extension("dyf.compact.tmp");
        let mut file = File::create(&tmp_path)?;

        // Header
        file.write_all(DYF2_MAGIC)?;
        file.write_all(&[0u8; 4])?;

        // Write all batches (merged base + overflow)
        let mut new_base_batches = Vec::new();
        let mut offset = 0u64;

        for (data, &rows) in all_batch_data.iter().zip(all_num_rows.iter()) {
            new_base_batches.push(BatchDescriptor {
                offset,
                length: data.len() as u64,
                num_rows: rows,
            });
            file.write_all(data)?;
            offset += data.len() as u64;
        }

        // Write field layer batches
        let mut new_field_layers = Vec::new();
        for (i, layer) in self.field_layers.iter().enumerate() {
            let mut new_batches = Vec::new();
            for (j, batch_data) in layer_data[i].iter().enumerate() {
                new_batches.push(BatchDescriptor {
                    offset,
                    length: batch_data.len() as u64,
                    num_rows: layer.batches[j].num_rows,
                });
                file.write_all(batch_data)?;
                offset += batch_data.len() as u64;
            }
            new_field_layers.push(FieldLayer {
                name: layer.name.clone(),
                fields: layer.fields.clone(),
                batches: new_batches,
            });
        }

        // Update state
        self.base_batches = new_base_batches;
        self.overflow_batches.clear();
        self.field_layers = new_field_layers;
        self.data_end = HEADER_SIZE + offset;

        // Write footer
        self.write_footer(&mut file)?;
        drop(file);

        // Replace original with compacted file
        std::fs::rename(&tmp_path, &self.path)?;

        Ok(())
    }

    /// Read a base batch's raw Arrow IPC bytes from the file.
    pub fn read_base_batch(&self, batch_idx: usize) -> io::Result<Vec<u8>> {
        self.read_batch_data(&self.base_batches[batch_idx])
    }

    /// Read an overflow batch's raw Arrow IPC bytes.
    pub fn read_overflow_batch(&self, batch_idx: usize) -> io::Result<Vec<u8>> {
        self.read_batch_data(&self.overflow_batches[batch_idx])
    }

    /// Read a field layer batch's raw Arrow IPC bytes.
    pub fn read_field_layer_batch(&self, layer_idx: usize, batch_idx: usize) -> io::Result<Vec<u8>> {
        self.read_batch_data(&self.field_layers[layer_idx].batches[batch_idx])
    }

    fn read_batch_data(&self, desc: &BatchDescriptor) -> io::Result<Vec<u8>> {
        let mut file = File::open(&self.path)?;
        file.seek(SeekFrom::Start(HEADER_SIZE + desc.offset))?;
        let mut buf = vec![0u8; desc.length as usize];
        file.read_exact(&mut buf)?;
        Ok(buf)
    }

    // Accessors
    pub fn path(&self) -> &Path {
        &self.path
    }
    pub fn tree(&self) -> &LshTree {
        &self.tree
    }
    pub fn tree_mut(&mut self) -> &mut LshTree {
        &mut self.tree
    }
    pub fn base_batches(&self) -> &[BatchDescriptor] {
        &self.base_batches
    }
    pub fn overflow_batches_desc(&self) -> &[BatchDescriptor] {
        &self.overflow_batches
    }
    pub fn field_layers(&self) -> &[FieldLayer] {
        &self.field_layers
    }
    pub fn metadata(&self) -> &HashMap<String, String> {
        &self.metadata
    }
    pub fn build_params(&self) -> &BuildParams {
        &self.build_params
    }
    pub fn embedding_dim(&self) -> u16 {
        self.embedding_dim
    }
    pub fn total_items(&self) -> u64 {
        self.total_items
    }

    /// Write the FlatBuffers index + 16-byte footer tail.
    fn write_footer(&self, file: &mut File) -> io::Result<()> {
        let fb_data = self.build_flatbuffer();
        let fb_offset = file.stream_position()?;

        file.write_all(&fb_data)?;

        // Write 16-byte footer: fb_offset (u64 LE) + fb_size (u64 LE)
        file.write_all(&fb_offset.to_le_bytes())?;
        file.write_all(&(fb_data.len() as u64).to_le_bytes())?;

        // Truncate file to current position (removes any old data beyond)
        let end = file.stream_position()?;
        file.set_len(end)?;

        Ok(())
    }

    /// Build the FlatBuffers index from current state.
    fn build_flatbuffer(&self) -> Vec<u8> {
        let mut fbb = flatbuffers::FlatBufferBuilder::with_capacity(4096);

        // Serialize nodes
        let nodes: Vec<_> = (0..self.tree.num_nodes())
            .map(|i| {
                let node = self.tree.node(i);
                serialize_node(&mut fbb, node)
            })
            .collect();
        let nodes_vec = fbb.create_vector(&nodes);

        // Serialize base batches
        let base_batches: Vec<_> = self
            .base_batches
            .iter()
            .map(|b| serialize_batch_desc(&mut fbb, b))
            .collect();
        let base_batches_vec = fbb.create_vector(&base_batches);

        // Serialize overflow batches
        let overflow_batches: Vec<_> = self
            .overflow_batches
            .iter()
            .map(|b| serialize_batch_desc(&mut fbb, b))
            .collect();
        let overflow_batches_vec = fbb.create_vector(&overflow_batches);

        // Serialize field layers
        let field_layers: Vec<_> = self
            .field_layers
            .iter()
            .map(|fl| serialize_field_layer(&mut fbb, fl))
            .collect();
        let field_layers_vec = fbb.create_vector(&field_layers);

        // Serialize build params
        let quant = fbb.create_string(&self.build_params.quantization);
        let comp = fbb.create_string(&self.build_params.compression);
        let bp = fb::BuildParams::create(
            &mut fbb,
            &fb::BuildParamsArgs {
                max_depth: self.build_params.max_depth,
                num_bits: self.build_params.num_bits,
                min_leaf_size: self.build_params.min_leaf_size,
                seed: self.build_params.seed,
                quantization: Some(quant),
                compression: Some(comp),
            },
        );

        // Serialize metadata
        let metadata: Vec<_> = self
            .metadata
            .iter()
            .map(|(k, v)| {
                let key = fbb.create_string(k);
                let val = fbb.create_string(v);
                fb::KeyValue::create(
                    &mut fbb,
                    &fb::KeyValueArgs {
                        key: Some(key),
                        value: Some(val),
                    },
                )
            })
            .collect();
        let metadata_vec = fbb.create_vector(&metadata);

        let version = fbb.create_string("2.0");

        let index = fb::Index::create(
            &mut fbb,
            &fb::IndexArgs {
                version: Some(version),
                embedding_dim: self.embedding_dim,
                total_items: self.total_items,
                num_leaves: self.tree.num_leaves() as u32,
                root: self.tree.root() as u32,
                nodes: Some(nodes_vec),
                batches: Some(base_batches_vec),
                overflow_batches: Some(overflow_batches_vec),
                field_layers: Some(field_layers_vec),
                build_params: Some(bp),
                metadata: Some(metadata_vec),
            },
        );

        fbb.finish(index, None);
        fbb.finished_data().to_vec()
    }
}

// --- Serialization helpers ---

fn serialize_node<'a, A: flatbuffers::Allocator + 'a>(
    fbb: &mut flatbuffers::FlatBufferBuilder<'a, A>,
    node: &TreeNode,
) -> flatbuffers::WIPOffset<fb::Node<'a>> {
    let children_vec = if node.children.is_empty() {
        None
    } else {
        let children: Vec<u32> = node.children.iter().map(|&c| c as u32).collect();
        Some(fbb.create_vector(&children))
    };

    let hp_vec = node.hyperplanes.as_ref().map(|hp| fbb.create_vector(hp));
    let ev_vec = node.eigenvalues.as_ref().map(|ev| fbb.create_vector(ev));
    let centroid_vec = fbb.create_vector(&node.centroid);

    let bucket_ids: Vec<u64> = node.bucket_id_to_child.iter().map(|&(bid, _)| bid).collect();
    let bucket_ids_vec = if bucket_ids.is_empty() {
        None
    } else {
        Some(fbb.create_vector(&bucket_ids))
    };

    fb::Node::create(
        fbb,
        &fb::NodeArgs {
            children: children_vec,
            hyperplanes: hp_vec,
            num_bits: node.num_bits,
            bucket_ids_to_children: bucket_ids_vec,
            centroid: Some(centroid_vec),
            num_items: node.num_items,
            batch_index: node.batch_index.map(|b| b as i32).unwrap_or(-1),
            depth: 0, // depth computed from tree structure
            eigenvalues: ev_vec,
        },
    )
}

fn serialize_batch_desc<'a, A: flatbuffers::Allocator + 'a>(
    fbb: &mut flatbuffers::FlatBufferBuilder<'a, A>,
    desc: &BatchDescriptor,
) -> flatbuffers::WIPOffset<fb::BatchDescriptor<'a>> {
    fb::BatchDescriptor::create(
        fbb,
        &fb::BatchDescriptorArgs {
            offset: desc.offset,
            length: desc.length,
            num_rows: desc.num_rows,
        },
    )
}

fn serialize_field_layer<'a, A: flatbuffers::Allocator + 'a>(
    fbb: &mut flatbuffers::FlatBufferBuilder<'a, A>,
    layer: &FieldLayer,
) -> flatbuffers::WIPOffset<fb::FieldLayer<'a>> {
    let name = fbb.create_string(&layer.name);

    let fields: Vec<_> = layer
        .fields
        .iter()
        .map(|fd| {
            let n = fbb.create_string(&fd.name);
            let t = fbb.create_string(&fd.arrow_type);
            fb::FieldDescriptor::create(
                fbb,
                &fb::FieldDescriptorArgs {
                    name: Some(n),
                    arrow_type: Some(t),
                },
            )
        })
        .collect();
    let fields_vec = fbb.create_vector(&fields);

    let batches: Vec<_> = layer
        .batches
        .iter()
        .map(|b| serialize_batch_desc(fbb, b))
        .collect();
    let batches_vec = fbb.create_vector(&batches);

    fb::FieldLayer::create(
        fbb,
        &fb::FieldLayerArgs {
            name: Some(name),
            fields: Some(fields_vec),
            batches: Some(batches_vec),
        },
    )
}

// --- Deserialization helpers ---

fn parse_tree_from_fb(
    index: &fb::Index,
    embedding_dim: usize,
    num_bits: usize,
    min_leaf_size: usize,
    seed: u64,
) -> LshTree {
    let fb_nodes = index.nodes().expect("Index must have nodes");
    let root = index.root() as usize;
    let max_leaf_size = min_leaf_size * 4; // default heuristic

    let mut tree = LshTree::new(embedding_dim, num_bits, min_leaf_size, max_leaf_size, seed);

    // Clear the default root node
    tree.clear_nodes();

    // Parse all nodes
    for i in 0..fb_nodes.len() {
        let fb_node = fb_nodes.get(i);

        let children: Vec<usize> = fb_node
            .children()
            .map(|c| (0..c.len()).map(|j| c.get(j) as usize).collect())
            .unwrap_or_default();

        let hyperplanes: Option<Vec<f32>> = fb_node.hyperplanes().map(|hp| {
            (0..hp.len()).map(|j| hp.get(j)).collect()
        });

        let eigenvalues: Option<Vec<f32>> = fb_node.eigenvalues().map(|ev| {
            (0..ev.len()).map(|j| ev.get(j)).collect()
        });

        let centroid: Vec<f32> = fb_node
            .centroid()
            .map(|c| (0..c.len()).map(|j| c.get(j)).collect())
            .unwrap_or_else(|| vec![0.0; embedding_dim]);

        let bucket_ids: Vec<u64> = fb_node
            .bucket_ids_to_children()
            .map(|b| (0..b.len()).map(|j| b.get(j)).collect())
            .unwrap_or_default();

        let bucket_id_to_child: Vec<(u64, usize)> = bucket_ids
            .iter()
            .zip(children.iter())
            .map(|(&bid, &child)| (bid, child))
            .collect();

        let is_leaf = children.is_empty();
        let batch_index = if fb_node.batch_index() >= 0 {
            Some(fb_node.batch_index() as u32)
        } else {
            None
        };

        let node = TreeNode {
            hyperplanes,
            eigenvalues,
            children,
            bucket_id_to_child,
            centroid,
            items: if is_leaf { Some(Vec::new()) } else { None }, // items loaded separately
            num_items: fb_node.num_items(),
            batch_index,
            parent: None, // reconstructed below
            num_bits: fb_node.num_bits(),
        };

        tree.push_node(node);
    }

    // Reconstruct parent pointers
    tree.set_root(root);
    tree.rebuild_parent_pointers();

    tree
}

fn parse_batch_descriptors(
    fb_batches: Option<flatbuffers::Vector<'_, flatbuffers::ForwardsUOffset<fb::BatchDescriptor<'_>>>>,
) -> Vec<BatchDescriptor> {
    fb_batches
        .map(|batches| {
            (0..batches.len())
                .map(|i| {
                    let b = batches.get(i);
                    BatchDescriptor {
                        offset: b.offset(),
                        length: b.length(),
                        num_rows: b.num_rows(),
                    }
                })
                .collect()
        })
        .unwrap_or_default()
}

fn parse_field_layers(
    fb_layers: Option<flatbuffers::Vector<'_, flatbuffers::ForwardsUOffset<fb::FieldLayer<'_>>>>,
) -> Vec<FieldLayer> {
    fb_layers
        .map(|layers| {
            (0..layers.len())
                .map(|i| {
                    let fl = layers.get(i);
                    let fields = fl
                        .fields()
                        .map(|fds| {
                            (0..fds.len())
                                .map(|j| {
                                    let fd = fds.get(j);
                                    FieldDescriptor {
                                        name: fd.name().unwrap_or("").to_string(),
                                        arrow_type: fd.arrow_type().unwrap_or("").to_string(),
                                    }
                                })
                                .collect()
                        })
                        .unwrap_or_default();

                    let batches = parse_batch_descriptors(fl.batches());

                    FieldLayer {
                        name: fl.name().unwrap_or("").to_string(),
                        fields,
                        batches,
                    }
                })
                .collect()
        })
        .unwrap_or_default()
}

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

    fn make_test_tree() -> LshTree {
        // Build a small tree from random data
        let dim = 8;
        let n = 32;
        let mut data = vec![0.0f32; n * dim];
        // Simple deterministic data
        for i in 0..n {
            for j in 0..dim {
                data[i * dim + j] = ((i * dim + j) as f32).sin();
            }
            // L2 normalize
            let norm: f32 = data[i * dim..(i + 1) * dim]
                .iter()
                .map(|x| x * x)
                .sum::<f32>()
                .sqrt();
            if norm > 0.0 {
                for j in 0..dim {
                    data[i * dim + j] /= norm;
                }
            }
        }
        LshTree::build(&data, n, dim, 2, 4, 16, 42)
    }

    #[test]
    fn test_create_and_open() {
        let tree = make_test_tree();
        let num_leaves = tree.num_leaves();
        let total = tree.total_items();

        // Create fake Arrow batches (just random bytes for testing)
        let leaf_indices = tree.leaf_indices();
        let leaf_batches: Vec<Vec<u8>> = leaf_indices
            .iter()
            .map(|&idx| {
                let n = tree.node(idx).num_items;
                vec![0xAB; (n as usize) * 10] // placeholder data
            })
            .collect();

        let bp = BuildParams {
            max_depth: 10,
            num_bits: 2,
            min_leaf_size: 4,
            seed: 42,
            quantization: "float16".into(),
            compression: "none".into(),
        };

        let mut metadata = HashMap::new();
        metadata.insert("test_key".into(), "test_value".into());

        let tmp = tempfile::NamedTempFile::new().unwrap();
        let path = tmp.path().to_str().unwrap();

        DyfFile::create(path, tree, &leaf_batches, bp, metadata).unwrap();

        // Open and verify
        let opened = DyfFile::open(path).unwrap();
        assert_eq!(opened.tree().num_leaves(), num_leaves);
        assert_eq!(opened.total_items(), total as u64);
        assert_eq!(opened.embedding_dim(), 8);
        assert_eq!(opened.metadata().get("test_key").unwrap(), "test_value");
        assert_eq!(opened.base_batches().len(), num_leaves);
        assert!(opened.overflow_batches_desc().is_empty());
        assert!(opened.field_layers().is_empty());
    }

    #[test]
    fn test_read_batch_roundtrip() {
        let tree = make_test_tree();
        let leaf_indices = tree.leaf_indices();
        let leaf_batches: Vec<Vec<u8>> = leaf_indices
            .iter()
            .enumerate()
            .map(|(i, &idx)| {
                let n = tree.node(idx).num_items;
                vec![(i as u8).wrapping_add(1); (n as usize) * 10]
            })
            .collect();

        let bp = BuildParams {
            max_depth: 10,
            num_bits: 2,
            min_leaf_size: 4,
            seed: 42,
            quantization: "float16".into(),
            compression: "none".into(),
        };

        let tmp = tempfile::NamedTempFile::new().unwrap();
        let path = tmp.path().to_str().unwrap();

        DyfFile::create(path, tree, &leaf_batches, bp, HashMap::new()).unwrap();

        let opened = DyfFile::open(path).unwrap();
        for (i, original) in leaf_batches.iter().enumerate() {
            let read_back = opened.read_base_batch(i).unwrap();
            assert_eq!(&read_back, original, "Batch {i} data mismatch");
        }
    }

    #[test]
    fn test_append_field_layer() {
        let tree = make_test_tree();
        let leaf_indices = tree.leaf_indices();
        let leaf_batches: Vec<Vec<u8>> = leaf_indices
            .iter()
            .map(|&idx| vec![0xFF; tree.node(idx).num_items as usize * 10])
            .collect();

        let bp = BuildParams {
            max_depth: 10,
            num_bits: 2,
            min_leaf_size: 4,
            seed: 42,
            quantization: "float16".into(),
            compression: "none".into(),
        };

        let tmp = tempfile::NamedTempFile::new().unwrap();
        let path = tmp.path().to_str().unwrap();

        let mut dyf =
            DyfFile::create(path, tree, &leaf_batches, bp, HashMap::new()).unwrap();

        // Append a field layer
        let field_data = vec![vec![0xCC; 100], vec![0xDD; 200]];
        let fields = vec![FieldDescriptor {
            name: "community_id".into(),
            arrow_type: "int32".into(),
        }];
        dyf.append_field_layer("community_id", fields, &field_data)
            .unwrap();

        // Verify by re-opening
        let reopened = DyfFile::open(path).unwrap();
        assert_eq!(reopened.field_layers().len(), 1);
        assert_eq!(reopened.field_layers()[0].name, "community_id");
        assert_eq!(reopened.field_layers()[0].batches.len(), 2);

        // Read back field layer batches
        let b0 = reopened.read_field_layer_batch(0, 0).unwrap();
        let b1 = reopened.read_field_layer_batch(0, 1).unwrap();
        assert_eq!(b0, vec![0xCC; 100]);
        assert_eq!(b1, vec![0xDD; 200]);

        // Original batches should still be readable
        for (i, original) in leaf_batches.iter().enumerate() {
            let read_back = reopened.read_base_batch(i).unwrap();
            assert_eq!(&read_back, original);
        }
    }

    #[test]
    fn test_append_overflow_items() {
        let tree = make_test_tree();
        let leaf_indices = tree.leaf_indices();
        let leaf_batches: Vec<Vec<u8>> = leaf_indices
            .iter()
            .map(|&idx| vec![0xAA; tree.node(idx).num_items as usize * 10])
            .collect();

        let bp = BuildParams {
            max_depth: 10,
            num_bits: 2,
            min_leaf_size: 4,
            seed: 42,
            quantization: "float16".into(),
            compression: "none".into(),
        };

        let tmp = tempfile::NamedTempFile::new().unwrap();
        let path = tmp.path().to_str().unwrap();

        let original_total = tree.total_items() as u64;
        let mut dyf =
            DyfFile::create(path, tree, &leaf_batches, bp, HashMap::new()).unwrap();

        // Append overflow items
        let overflow = vec![vec![0xBB; 500]];
        dyf.append_items(&overflow, &[50]).unwrap();

        assert_eq!(dyf.total_items(), original_total + 50);
        assert!(dyf.fragmentation() > 0.0);

        // Re-open and verify
        let reopened = DyfFile::open(path).unwrap();
        assert_eq!(reopened.total_items(), original_total + 50);
        assert_eq!(reopened.overflow_batches_desc().len(), 1);

        let ob = reopened.read_overflow_batch(0).unwrap();
        assert_eq!(ob, vec![0xBB; 500]);
    }

    #[test]
    fn test_update_metadata() {
        let tree = make_test_tree();
        let leaf_indices = tree.leaf_indices();
        let leaf_batches: Vec<Vec<u8>> = leaf_indices
            .iter()
            .map(|_| vec![0x00; 10])
            .collect();

        let bp = BuildParams {
            max_depth: 10,
            num_bits: 2,
            min_leaf_size: 4,
            seed: 42,
            quantization: "float16".into(),
            compression: "none".into(),
        };

        let tmp = tempfile::NamedTempFile::new().unwrap();
        let path = tmp.path().to_str().unwrap();

        let mut dyf =
            DyfFile::create(path, tree, &leaf_batches, bp, HashMap::new()).unwrap();

        dyf.update_metadata(&[("key1", "value1"), ("key2", "value2")])
            .unwrap();

        let reopened = DyfFile::open(path).unwrap();
        assert_eq!(reopened.metadata().get("key1").unwrap(), "value1");
        assert_eq!(reopened.metadata().get("key2").unwrap(), "value2");
    }

    #[test]
    fn test_compact() {
        let tree = make_test_tree();
        let leaf_indices = tree.leaf_indices();
        let original_base_count = leaf_indices.len();
        let leaf_batches: Vec<Vec<u8>> = leaf_indices
            .iter()
            .map(|&idx| vec![0xAA; tree.node(idx).num_items as usize * 10])
            .collect();

        let bp = BuildParams {
            max_depth: 10,
            num_bits: 2,
            min_leaf_size: 4,
            seed: 42,
            quantization: "float16".into(),
            compression: "none".into(),
        };

        let tmp = tempfile::NamedTempFile::new().unwrap();
        let path = tmp.path().to_str().unwrap();

        let original_total = tree.total_items() as u64;
        let mut dyf =
            DyfFile::create(path, tree, &leaf_batches, bp, HashMap::new()).unwrap();

        // Append overflow
        let overflow = vec![vec![0xBB; 500]];
        dyf.append_items(&overflow, &[50]).unwrap();
        assert!(dyf.fragmentation() > 0.0);
        assert_eq!(dyf.overflow_batches_desc().len(), 1);

        // Also append a field layer
        let field_data = vec![vec![0xCC; 100]];
        dyf.append_field_layer(
            "test_field",
            vec![FieldDescriptor { name: "test".into(), arrow_type: "int32".into() }],
            &field_data,
        ).unwrap();

        // Compact
        dyf.compact().unwrap();

        // Overflow should be empty, merged into base
        assert_eq!(dyf.fragmentation(), 0.0);
        assert!(dyf.overflow_batches_desc().is_empty());
        assert_eq!(
            dyf.base_batches().len(),
            original_base_count + 1, // merged overflow
        );
        assert_eq!(dyf.total_items(), original_total + 50);

        // Field layers should be preserved
        assert_eq!(dyf.field_layers().len(), 1);

        // Verify by re-opening
        let reopened = DyfFile::open(path).unwrap();
        assert_eq!(reopened.total_items(), original_total + 50);
        assert!(reopened.overflow_batches_desc().is_empty());
        assert_eq!(reopened.base_batches().len(), original_base_count + 1);
        assert_eq!(reopened.field_layers().len(), 1);

        // Read back the field layer data
        let fl_data = reopened.read_field_layer_batch(0, 0).unwrap();
        assert_eq!(fl_data, vec![0xCC; 100]);
    }
}
