use std::{collections::HashMap, hash::Hash};

use pumpkin_data::{
    BlockState, BlockStateId,
    block_properties::{has_random_ticks, is_air, is_liquid},
    fluid::Fluid,
};
use pumpkin_util::encompassing_bits;
use tracing::warn;

use super::format::{ChunkSectionBiomes, ChunkSectionBlockStates};

/// 3d array indexed by y,z,x
type AbstractCube<T, const DIM: usize> = [[[T; DIM]; DIM]; DIM];

#[inline]
#[must_use]
pub fn has_random_ticking_fluid(id: BlockStateId) -> bool {
    Fluid::from_state_id(id).is_some_and(|fluid| Fluid::same_fluid_type(fluid.id, Fluid::LAVA.id))
}

#[derive(Clone)]
pub struct HeterogeneousPaletteData<V: Hash + Eq + Copy, const DIM: usize> {
    storage: PaletteStorage<V, DIM>,
    palette: Vec<V>,
    counts: Vec<u16>,
}

#[derive(Clone)]
enum PaletteStorage<V, const DIM: usize> {
    Dense(Box<AbstractCube<V, DIM>>),
    Indexed(Box<AbstractCube<u8, DIM>>),
}

impl<V: Hash + Eq + Copy + Default, const DIM: usize> HeterogeneousPaletteData<V, DIM> {
    fn get(&self, x: usize, y: usize, z: usize) -> V {
        debug_assert!(x < DIM);
        debug_assert!(y < DIM);
        debug_assert!(z < DIM);

        match &self.storage {
            PaletteStorage::Dense(cube) => cube[y][z][x],
            PaletteStorage::Indexed(indices) => self.palette[indices[y][z][x] as usize],
        }
    }

    /// Returns the Original
    fn set(&mut self, x: usize, y: usize, z: usize, value: V) -> V {
        debug_assert!(x < DIM);
        debug_assert!(y < DIM);
        debug_assert!(z < DIM);

        let original = self.get(x, y, z);
        if original == value {
            return original;
        }

        let original_index = self.palette.iter().position(|v| v == &original).unwrap();

        // Find or add the new value to the palette.
        let new_index = if let Some(new_index) = self.palette.iter().position(|v| v == &value) {
            self.counts[new_index] += 1;
            new_index
        } else {
            self.palette.push(value);
            self.counts.push(1);
            self.palette.len() - 1
        };

        // Handle storage upgrades or updates
        let mut upgraded = false;
        match &mut self.storage {
            PaletteStorage::Dense(cube) => {
                cube[y][z][x] = value;
            }
            PaletteStorage::Indexed(indices) => {
                if new_index <= 255 {
                    indices[y][z][x] = new_index as u8;
                } else {
                    // Upgrade to Dense
                    let mut cube = Box::new([[[V::default(); DIM]; DIM]; DIM]);
                    for (i, v) in cube
                        .as_flattened_mut()
                        .as_flattened_mut()
                        .iter_mut()
                        .enumerate()
                    {
                        let y = i / (DIM * DIM);
                        let z = (i / DIM) % DIM;
                        let x = i % DIM;
                        *v = self.palette[indices[y][z][x] as usize];
                    }
                    cube[y][z][x] = value;
                    self.storage = PaletteStorage::Dense(cube);
                    upgraded = true;
                }
            }
        }

        self.counts[original_index] -= 1;

        if self.counts[original_index] == 0 {
            let last_index = self.palette.len() - 1;
            // Remove from palette and counts Vecs if the count hits zero.
            self.palette.swap_remove(original_index);
            self.counts.swap_remove(original_index);

            if self.palette.capacity() > 16 && self.palette.len() < self.palette.capacity() / 2 {
                self.palette.shrink_to_fit();
                self.counts.shrink_to_fit();
            }

            // If we are indexed, we need to update all indices because swap_remove changed indices
            if !upgraded && let PaletteStorage::Indexed(indices) = &mut self.storage {
                for row in indices.iter_mut() {
                    for col in row.iter_mut() {
                        for idx in col.iter_mut() {
                            if *idx as usize == last_index {
                                *idx = original_index as u8;
                            }
                        }
                    }
                }
            }
        }

        original
    }
}

/// A paletted container is a cube of registry ids. It uses a custom compression scheme based on how
/// may distinct registry ids are in the cube.
#[derive(Clone)]
pub enum PalettedContainer<V: Hash + Eq + Copy + Default, const DIM: usize> {
    Homogeneous(V),
    Heterogeneous(Box<HeterogeneousPaletteData<V, DIM>>),
}

impl<V: Hash + Eq + Copy + Default, const DIM: usize> PalettedContainer<V, DIM> {
    pub const SIZE: usize = DIM;
    pub const VOLUME: usize = DIM * DIM * DIM;

    fn from_cube(cube: Box<AbstractCube<V, DIM>>) -> Self {
        let mut palette: Vec<V> = Vec::new();
        let mut counts: Vec<u16> = Vec::new();

        // Iterate over the flattened cube to populate the palette and counts
        for val in cube.as_flattened().as_flattened() {
            if let Some(index) = palette.iter().position(|v| v == val) {
                // Value already exists, increment its count
                counts[index] += 1;
            } else {
                // New value, add it to the palette and start its count
                palette.push(*val);
                counts.push(1);
            }
        }

        if palette.len() == 1 {
            // Fast path: the cube is homogeneous, so we can store just one value
            Self::Homogeneous(palette[0])
        } else {
            // Heterogeneous cube, store the full data
            if palette.len() <= 256 && std::mem::size_of::<V>() > 1 {
                let mut indices = Box::new([[[0u8; DIM]; DIM]; DIM]);
                for (i, v) in cube.as_flattened().as_flattened().iter().enumerate() {
                    let idx = palette.iter().position(|p| p == v).unwrap();
                    indices.as_flattened_mut().as_flattened_mut()[i] = idx as u8;
                }
                Self::Heterogeneous(Box::new(HeterogeneousPaletteData {
                    storage: PaletteStorage::Indexed(indices),
                    palette,
                    counts,
                }))
            } else {
                Self::Heterogeneous(Box::new(HeterogeneousPaletteData {
                    storage: PaletteStorage::Dense(cube),
                    palette,
                    counts,
                }))
            }
        }
    }

    fn bits_per_entry(&self) -> u8 {
        match self {
            Self::Homogeneous(_) => 0,
            Self::Heterogeneous(data) => encompassing_bits(data.counts.len()),
        }
    }

    pub fn to_palette_and_packed_data(&self, bits_per_entry: u8) -> (Box<[V]>, Box<[i64]>) {
        match self {
            Self::Homogeneous(registry_id) => (Box::new([*registry_id]), Box::new([])),
            Self::Heterogeneous(data) => {
                debug_assert!(bits_per_entry >= encompassing_bits(data.counts.len()));
                debug_assert!(bits_per_entry <= 15);

                // Don't use HashMap's here, because its slow
                let blocks_per_i64 = 64 / bits_per_entry;

                let packed_indices: Box<[i64]> = match &data.storage {
                    PaletteStorage::Dense(cube) => cube
                        .as_flattened()
                        .as_flattened()
                        .chunks(blocks_per_i64 as usize)
                        .map(|chunk| {
                            chunk.iter().enumerate().fold(0, |acc, (index, key)| {
                                let key_index =
                                    data.palette.iter().position(|&x| x == *key).unwrap();
                                debug_assert!((1 << bits_per_entry) > key_index);

                                let packed_offset_index =
                                    (key_index as u64) << (bits_per_entry as u64 * index as u64);
                                acc | packed_offset_index as i64
                            })
                        })
                        .collect(),
                    PaletteStorage::Indexed(indices) => indices
                        .as_flattened()
                        .as_flattened()
                        .chunks(blocks_per_i64 as usize)
                        .map(|chunk| {
                            chunk.iter().enumerate().fold(0, |acc, (index, key_index)| {
                                let key_index = *key_index as usize;
                                debug_assert!((1 << bits_per_entry) > key_index);

                                let packed_offset_index =
                                    (key_index as u64) << (bits_per_entry as u64 * index as u64);
                                acc | packed_offset_index as i64
                            })
                        })
                        .collect(),
                };

                (data.palette.clone().into_boxed_slice(), packed_indices)
            }
        }
    }

    #[must_use]
    pub fn from_palette_and_packed_data(
        palette: &[V],
        packed_data: &[i64],
        minimum_bits_per_entry: u8,
    ) -> Self {
        if palette.is_empty() {
            warn!("No palette data! Defaulting...");
            return Self::Homogeneous(V::default());
        }

        if palette.len() == 1 {
            return Self::Homogeneous(palette[0]);
        }

        let bits_per_key = encompassing_bits(palette.len()).max(minimum_bits_per_entry);
        let index_mask = (1 << bits_per_key) - 1;
        let keys_per_i64 = 64 / bits_per_key;

        // Optimized path for indexed storage if palette is small enough
        if palette.len() <= 256 && std::mem::size_of::<V>() > 1 {
            let mut indices = Box::new([[[0u8; DIM]; DIM]; DIM]);
            let mut counts = vec![0u16; palette.len()];
            let indices_flat = indices.as_flattened_mut().as_flattened_mut();

            let mut packed_data_iter = packed_data.iter();
            let mut current_packed_word = *packed_data_iter.next().unwrap_or(&0);

            for (i, index_out) in indices_flat.iter_mut().enumerate().take(Self::VOLUME) {
                let bit_index_in_word = i % keys_per_i64 as usize;
                if bit_index_in_word == 0 && i > 0 {
                    current_packed_word = *packed_data_iter.next().unwrap_or(&0);
                }

                let lookup_index = ((current_packed_word as u64)
                    >> (bit_index_in_word as u64 * bits_per_key as u64))
                    & index_mask;

                let idx = lookup_index as usize;
                if idx < palette.len() {
                    *index_out = idx as u8;
                    counts[idx] += 1;
                } else {
                    warn!("Lookup index out of bounds! Defaulting...");
                    // value is already 0, and counts[0] will be updated correctly if we track it
                }
            }
            // fix counts[0] if it was skipped in out-of-bounds cases (rare)
            // But actually we should just ensure it's correct.

            return Self::Heterogeneous(Box::new(HeterogeneousPaletteData {
                storage: PaletteStorage::Indexed(indices),
                palette: palette.to_vec(),
                counts,
            }));
        }

        let mut decompressed_values = Vec::with_capacity(Self::VOLUME);

        let mut packed_data_iter = packed_data.iter();
        let mut current_packed_word = *packed_data_iter.next().unwrap_or(&0);

        for i in 0..Self::VOLUME {
            let bit_index_in_word = i % keys_per_i64 as usize;

            if bit_index_in_word == 0 && i > 0 {
                current_packed_word = *packed_data_iter.next().unwrap_or(&0);
            }

            let lookup_index = (current_packed_word as u64
                >> (bit_index_in_word as u64 * bits_per_key as u64))
                & index_mask;

            let value = palette
                .get(lookup_index as usize)
                .copied()
                .unwrap_or_else(|| {
                    warn!("Lookup index out of bounds! Defaulting...");
                    V::default()
                });

            decompressed_values.push(value);
        }

        // Now, with all decompressed values, build the counts.
        let mut counts = vec![0; palette.len()];

        for &value in &decompressed_values {
            // This is the key optimization: find the index in the palette Vec
            // and increment the corresponding count.
            if let Some(index) = palette.iter().position(|v| v == &value) {
                counts[index] += 1;
            } else {
                // This case should ideally not happen if the palette is complete.
                warn!("Decompressed value not found in palette!");
            }
        }

        let mut cube = Box::new([[[V::default(); DIM]; DIM]; DIM]);
        cube.as_flattened_mut()
            .as_flattened_mut()
            .copy_from_slice(&decompressed_values);

        Self::Heterogeneous(Box::new(HeterogeneousPaletteData {
            storage: PaletteStorage::Dense(cube),
            palette: palette.to_vec(),
            counts,
        }))
    }

    pub fn get(&self, x: usize, y: usize, z: usize) -> V {
        match self {
            Self::Homogeneous(value) => *value,
            Self::Heterogeneous(data) => data.get(x, y, z),
        }
    }

    pub fn set(&mut self, x: usize, y: usize, z: usize, value: V) -> V {
        debug_assert!(x < Self::SIZE);
        debug_assert!(y < Self::SIZE);
        debug_assert!(z < Self::SIZE);

        match self {
            Self::Homogeneous(original) => {
                let original = *original;
                if value != original {
                    let mut cube = Box::new([[[original; DIM]; DIM]; DIM]);
                    cube[y][z][x] = value;
                    *self = Self::from_cube(cube);
                }
                original
            }
            Self::Heterogeneous(data) => {
                let original = data.set(x, y, z, value);
                if data.counts.len() == 1 {
                    *self = Self::Homogeneous(data.palette[0]);
                }
                original
            }
        }
    }

    pub fn iter(&self) -> Box<dyn Iterator<Item = V> + '_> {
        match self {
            Self::Homogeneous(registry_id) => {
                Box::new(std::iter::repeat_n(*registry_id, Self::VOLUME))
            }
            Self::Heterogeneous(data) => match &data.storage {
                PaletteStorage::Dense(cube) => {
                    Box::new(cube.as_flattened().as_flattened().iter().copied())
                }
                PaletteStorage::Indexed(indices) => Box::new(
                    indices
                        .as_flattened()
                        .as_flattened()
                        .iter()
                        .map(|&idx| data.palette[idx as usize]),
                ),
            },
        }
    }

    pub fn is_empty(&self) -> bool {
        match self {
            Self::Homogeneous(value) => *value == V::default(),
            Self::Heterogeneous(_) => false,
        }
    }
}

impl<'a, V: Hash + Eq + Copy + Default, const DIM: usize> IntoIterator
    for &'a PalettedContainer<V, DIM>
{
    type Item = V;
    type IntoIter = Box<dyn Iterator<Item = Self::Item> + 'a>;

    fn into_iter(self) -> Self::IntoIter {
        match self {
            PalettedContainer::Homogeneous(registry_id) => Box::new(std::iter::repeat_n(
                *registry_id,
                PalettedContainer::<V, DIM>::VOLUME,
            )),
            PalettedContainer::Heterogeneous(data) => match &data.storage {
                PaletteStorage::Dense(cube) => {
                    Box::new(cube.as_flattened().as_flattened().iter().copied())
                }
                PaletteStorage::Indexed(indices) => Box::new(
                    indices
                        .as_flattened()
                        .as_flattened()
                        .iter()
                        .map(|&idx| data.palette[idx as usize]),
                ),
            },
        }
    }
}

impl<V: Default + Hash + Eq + Copy, const DIM: usize> Default for PalettedContainer<V, DIM> {
    fn default() -> Self {
        Self::Homogeneous(V::default())
    }
}

impl BiomePalette {
    #[must_use]
    pub fn convert_network(&self) -> NetworkSerialization<u8> {
        match self {
            Self::Homogeneous(registry_id) => NetworkSerialization {
                bits_per_entry: 0,
                palette: NetworkPalette::Single(*registry_id),
                packed_data: Box::new([]),
            },
            Self::Heterogeneous(data) => {
                let raw_bits_per_entry = encompassing_bits(data.counts.len());
                if raw_bits_per_entry > BIOME_NETWORK_MAX_MAP_BITS {
                    let bits_per_entry = BIOME_NETWORK_MAX_BITS;
                    let values_per_i64 = 64 / bits_per_entry;
                    let mut packed_data = Vec::new();
                    let mut current_idx = 0;
                    while current_idx < Self::VOLUME {
                        let mut acc = 0u64;
                        for i in 0..values_per_i64 as usize {
                            if current_idx + i < Self::VOLUME {
                                let y = (current_idx + i) / (Self::SIZE * Self::SIZE);
                                let z = ((current_idx + i) / Self::SIZE) % Self::SIZE;
                                let x = (current_idx + i) % Self::SIZE;
                                let value = data.get(x, y, z);
                                debug_assert!((1 << bits_per_entry) > value);
                                acc |= (value as u64) << (bits_per_entry as u64 * i as u64);
                            }
                        }
                        packed_data.push(acc as i64);
                        current_idx += values_per_i64 as usize;
                    }

                    NetworkSerialization {
                        bits_per_entry,
                        palette: NetworkPalette::Direct,
                        packed_data: packed_data.into_boxed_slice(),
                    }
                } else {
                    let bits_per_entry = raw_bits_per_entry.max(BIOME_NETWORK_MIN_MAP_BITS);
                    let (palette, packed) = self.to_palette_and_packed_data(bits_per_entry);

                    NetworkSerialization {
                        bits_per_entry,
                        palette: NetworkPalette::Indirect(palette),
                        packed_data: packed,
                    }
                }
            }
        }
    }

    #[must_use]
    pub fn convert_be_network(&self) -> BeNetworkSerialization<u8> {
        match self {
            Self::Homogeneous(registry_id) => BeNetworkSerialization {
                bits_per_entry: 0,
                palette: NetworkPalette::Single(
                    pumpkin_data::biome::Biome::from_id(*registry_id)
                        .map_or(1, |b| b.be_network_id),
                ),
                packed_data: Box::new([]),
            },
            Self::Heterogeneous(data) => {
                let bits_per_entry = encompassing_bits(data.palette.len());

                let key_to_index_map: HashMap<_, usize> = data
                    .palette
                    .iter()
                    .enumerate()
                    .map(|(index, key)| (*key, index))
                    .collect();

                let biomes_per_word = 32 / bits_per_entry;
                let volume: usize = 4096;
                let expected_word_count = volume.div_ceil(biomes_per_word as usize);
                let mut packed_data = Vec::with_capacity(expected_word_count);

                let mut current_word: u32 = 0;
                let mut current_index_in_word = 0;

                for x in 0..16 {
                    for y in 0..16 {
                        for z in 0..16 {
                            let key = self.get(x / 4, z / 4, y / 4);
                            let key_index = key_to_index_map.get(&key).unwrap();
                            debug_assert!((1 << bits_per_entry) > *key_index);

                            current_word |= (*key_index as u32)
                                << (bits_per_entry as u32 * current_index_in_word);
                            current_index_in_word += 1;

                            if current_index_in_word == biomes_per_word as u32 {
                                packed_data.push(current_word);
                                current_word = 0;
                                current_index_in_word = 0;
                            }
                        }
                    }
                }

                // Push any remaining bits if the volume isn't a multiple of biomes_per_word
                if current_index_in_word > 0 {
                    packed_data.push(current_word);
                }

                BeNetworkSerialization {
                    bits_per_entry,
                    palette: NetworkPalette::Indirect(
                        data.palette
                            .iter()
                            .map(|&id| {
                                pumpkin_data::biome::Biome::from_id(id)
                                    .map_or(1, |b| b.be_network_id)
                            })
                            .collect::<Vec<_>>()
                            .into_boxed_slice(),
                    ),
                    packed_data: packed_data.into_boxed_slice(),
                }
            }
        }
    }

    #[must_use]
    pub fn from_disk_nbt(nbt: ChunkSectionBiomes) -> Self {
        let palette = nbt.palette;

        Self::from_palette_and_packed_data(
            &palette,
            nbt.data.as_ref().unwrap_or(&Box::default()),
            BIOME_DISK_MIN_BITS,
        )
    }

    #[must_use]
    pub fn to_disk_nbt(&self) -> ChunkSectionBiomes {
        #[expect(clippy::unnecessary_min_or_max)]
        let bits_per_entry = self.bits_per_entry().max(BIOME_DISK_MIN_BITS);
        let (palette, packed_data) = self.to_palette_and_packed_data(bits_per_entry);
        ChunkSectionBiomes {
            data: if packed_data.is_empty() {
                None
            } else {
                Some(packed_data)
            },
            palette,
        }
    }
}

impl BlockPalette {
    #[must_use]
    pub fn convert_network(&self) -> NetworkSerialization<u16> {
        match self {
            Self::Homogeneous(registry_id) => NetworkSerialization {
                bits_per_entry: 0,
                palette: NetworkPalette::Single(registry_id.as_u16()),
                packed_data: Box::new([]),
            },
            Self::Heterogeneous(data) => {
                let raw_bits_per_entry = encompassing_bits(data.counts.len());
                if raw_bits_per_entry > BLOCK_NETWORK_MAX_MAP_BITS {
                    let bits_per_entry = BLOCK_NETWORK_MAX_BITS;
                    let values_per_i64 = 64 / bits_per_entry;
                    let mut packed_data = Vec::new();
                    let mut current_idx = 0;
                    while current_idx < Self::VOLUME {
                        let mut acc = 0u64;
                        for i in 0..values_per_i64 as usize {
                            if current_idx + i < Self::VOLUME {
                                let y = (current_idx + i) / (Self::SIZE * Self::SIZE);
                                let z = ((current_idx + i) / Self::SIZE) % Self::SIZE;
                                let x = (current_idx + i) % Self::SIZE;
                                let value = data.get(x, y, z).as_u16();
                                debug_assert!((1 << bits_per_entry) > value);
                                acc |= (value as u64) << (bits_per_entry as u64 * i as u64);
                            }
                        }
                        packed_data.push(acc as i64);
                        current_idx += values_per_i64 as usize;
                    }

                    NetworkSerialization {
                        bits_per_entry,
                        palette: NetworkPalette::Direct,
                        packed_data: packed_data.into_boxed_slice(),
                    }
                } else {
                    let bits_per_entry = raw_bits_per_entry.max(BLOCK_NETWORK_MIN_MAP_BITS);
                    let (palette, packed) = self.to_palette_and_packed_data(bits_per_entry);

                    NetworkSerialization {
                        bits_per_entry,
                        palette: NetworkPalette::Indirect(
                            palette.iter().map(|v| v.as_u16()).collect(),
                        ),
                        packed_data: packed,
                    }
                }
            }
        }
    }

    #[must_use]
    pub fn convert_be_network(&self) -> BeNetworkSerialization<u16> {
        match self {
            Self::Homogeneous(registry_id) => BeNetworkSerialization {
                bits_per_entry: 0,
                palette: NetworkPalette::Single(BlockState::to_be_network_id(*registry_id)),
                packed_data: Box::new([]),
            },
            Self::Heterogeneous(data) => {
                let bits_per_entry = encompassing_bits(data.palette.len());

                let key_to_index_map: HashMap<_, usize> = data
                    .palette
                    .iter()
                    .enumerate()
                    .map(|(index, key)| (*key, index))
                    .collect();

                let blocks_per_word = 32 / bits_per_entry;
                let expected_word_count = Self::VOLUME.div_ceil(blocks_per_word as usize);
                let mut packed_data = Vec::with_capacity(expected_word_count);

                let mut current_word: u32 = 0;
                let mut current_index_in_word = 0;

                for x in 0..16 {
                    for y in 0..16 {
                        for z in 0..16 {
                            // Java has it in y, z, x order, so we need to convert it back to x, y, z
                            // Please test your code on bedrock before merging
                            let key = data.get(x, z, y);
                            let key_index = key_to_index_map.get(&key).unwrap();
                            debug_assert!((1 << bits_per_entry) > *key_index);

                            current_word |= (*key_index as u32)
                                << (bits_per_entry as u32 * current_index_in_word);
                            current_index_in_word += 1;

                            if current_index_in_word == blocks_per_word as u32 {
                                packed_data.push(current_word);
                                current_word = 0;
                                current_index_in_word = 0;
                            }
                        }
                    }
                }

                // Push any remaining bits if the volume isn't a multiple of blocks_per_word
                if current_index_in_word > 0 {
                    packed_data.push(current_word);
                }

                BeNetworkSerialization {
                    bits_per_entry,
                    palette: NetworkPalette::Indirect(
                        data.palette
                            .iter()
                            .map(|&id| BlockState::to_be_network_id(id))
                            .collect(),
                    ),
                    packed_data: packed_data.into_boxed_slice(),
                }
            }
        }
    }

    /// Check if the entire chunk is filled with only air
    #[must_use]
    pub fn has_only_air(&self) -> bool {
        match self {
            Self::Homogeneous(id) => is_air(*id),
            Self::Heterogeneous(data) => data.palette.iter().all(|&id| is_air(id)),
        }
    }

    #[must_use]
    pub fn random_ticking_counts(&self) -> (u16, u16) {
        match self {
            Self::Homogeneous(registry_id) => {
                let block_count = if has_random_ticks(*registry_id) {
                    Self::VOLUME as u16
                } else {
                    0
                };
                let fluid_count = if has_random_ticking_fluid(*registry_id) {
                    Self::VOLUME as u16
                } else {
                    0
                };
                (block_count, fluid_count)
            }
            Self::Heterogeneous(data) => data.palette.iter().zip(data.counts.iter()).fold(
                (0, 0),
                |(block_count, fluid_count), (registry_id, count)| {
                    let block_count = if has_random_ticks(*registry_id) {
                        block_count.saturating_add(*count)
                    } else {
                        block_count
                    };
                    let fluid_count = if has_random_ticking_fluid(*registry_id) {
                        fluid_count.saturating_add(*count)
                    } else {
                        fluid_count
                    };
                    (block_count, fluid_count)
                },
            ),
        }
    }

    #[must_use]
    pub fn non_air_block_count(&self) -> u16 {
        match self {
            Self::Homogeneous(registry_id) => {
                if is_air(*registry_id) {
                    0
                } else {
                    Self::VOLUME as u16
                }
            }
            Self::Heterogeneous(data) => data
                .palette
                .iter()
                .zip(data.counts.iter())
                .filter_map(|(registry_id, count)| (!is_air(*registry_id)).then_some(*count))
                .sum(),
        }
    }

    #[must_use]
    pub fn liquid_block_count(&self) -> u16 {
        match self {
            Self::Homogeneous(registry_id) => {
                if is_liquid(*registry_id) {
                    Self::VOLUME as u16
                } else {
                    0
                }
            }
            Self::Heterogeneous(data) => data
                .palette
                .iter()
                .zip(data.counts.iter())
                .filter_map(|(registry_id, count)| is_liquid(*registry_id).then_some(*count))
                .sum(),
        }
    }

    #[must_use]
    pub fn from_disk_nbt(nbt: ChunkSectionBlockStates) -> Self {
        let palette = nbt.palette;

        Self::from_palette_and_packed_data(
            &palette,
            nbt.data.as_ref().unwrap_or(&Box::default()),
            BLOCK_DISK_MIN_BITS,
        )
    }

    #[must_use]
    pub fn to_disk_nbt(&self) -> ChunkSectionBlockStates {
        let bits_per_entry = self.bits_per_entry().max(BLOCK_DISK_MIN_BITS);
        let (palette, packed_data) = self.to_palette_and_packed_data(bits_per_entry);
        ChunkSectionBlockStates {
            data: if packed_data.is_empty() {
                None
            } else {
                Some(packed_data)
            },
            palette,
        }
    }
}

/// Represents the different types of data encoding used in Minecraft's bit-packed chunk sections.
///
/// Minecraft uses a "Palette" system to compress chunk data. Instead of sending a full
/// 15-bit `BlockState` ID for every block, it sends a smaller index (e.g., 4 bits) that
/// points to a value in these palettes.
pub enum NetworkPalette<V> {
    /// **Single Value Palette (Bits per entry: 0)**
    ///
    /// Used when an entire chunk section (16x16x16) consists of only one type of block or biome.
    /// No data array follows this palette in the network buffer.
    Single(V),
    /// **Indirect Palette (Bits per entry: 1-8 for blocks, 1-3 for biomes)**
    ///
    /// A list of unique values present in the section. The data array contains indices
    /// pointing into this list.
    Indirect(Box<[V]>),
    /// **Direct Palette (Bits per entry: 15+ for blocks, 6+ for biomes)**
    ///
    /// Used when the section is too complex for a small palette. The data array
    /// contains global Registry IDs directly. No palette list is sent.
    Direct,
}

pub struct NetworkSerialization<V> {
    pub bits_per_entry: u8,
    pub palette: NetworkPalette<V>,
    pub packed_data: Box<[i64]>,
}

pub struct BeNetworkSerialization<V> {
    pub bits_per_entry: u8,
    pub palette: NetworkPalette<V>,
    pub packed_data: Box<[u32]>,
}

// According to the wiki, palette serialization for disk and network is different. Disk
// serialization always uses a palette if greater than one entry. Network serialization packs ids
// directly instead of using a palette above a certain bits-per-entry

// TODO: Do our own testing; do we really need to handle network and disk serialization differently?
pub type BlockPalette = PalettedContainer<BlockStateId, 16>;
const BLOCK_DISK_MIN_BITS: u8 = 4;
const BLOCK_NETWORK_MIN_MAP_BITS: u8 = 4;
const BLOCK_NETWORK_MAX_MAP_BITS: u8 = 8;
pub(crate) const BLOCK_NETWORK_MAX_BITS: u8 = 15;

pub type BiomePalette = PalettedContainer<u8, 4>;
const BIOME_DISK_MIN_BITS: u8 = 0;
const BIOME_NETWORK_MIN_MAP_BITS: u8 = 1;
const BIOME_NETWORK_MAX_MAP_BITS: u8 = 3;
pub(crate) const BIOME_NETWORK_MAX_BITS: u8 = 7;