// SPDX-License-Identifier: AGPL-3.0-only

//! Dense and palette-compressed chunk storage forms.

use super::{BlockId, CHUNK_SIZE, CHUNK_VOLUME};
use serde::{Deserialize, Serialize};
use std::collections::{HashMap, hash_map::Entry};

/// A spatial volume containing voxel data.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct Chunk {
    /// The array of blocks contained within the chunk.
    pub blocks: Box<[BlockId]>,
}

impl Chunk {
    /// Calculates the one-dimensional index for a three-dimensional block coordinate.
    #[inline]
    #[must_use]
    pub const fn index(x: usize, y: usize, z: usize) -> usize {
        (x * CHUNK_SIZE * CHUNK_SIZE) + (z * CHUNK_SIZE) + y
    }

    /// Retrieves the block identifier at the specified coordinates.
    #[must_use]
    pub fn get(&self, x: usize, y: usize, z: usize) -> BlockId {
        self.blocks[Self::index(x, y, z)]
    }

    /// Sets the block identifier at the specified coordinates.
    pub fn set(&mut self, x: usize, y: usize, z: usize, block: BlockId) {
        self.blocks[Self::index(x, y, z)] = block;
    }
}

impl Default for Chunk {
    fn default() -> Self {
        Self {
            blocks: vec![BlockId::AIR; CHUNK_VOLUME].into_boxed_slice(),
        }
    }
}

/// The palette-compressed representation of a chunk's voxel data, used as the stored and transmitted form.
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct PalettedChunk {
    /// The distinct materials present in the chunk. A stored voxel index refers to a position in this table.
    palette: Vec<BlockId>,
    /// The palette indices for all `CHUNK_VOLUME` voxels, packed `bits_per_index` bits each, low voxel first, into 64-bit words.
    indices: Vec<u64>,
    /// The width in bits of a single packed palette index, equal to `ceil(log2(palette.len()))` with a floor of one.
    bits_per_index: u32,
}

impl PalettedChunk {
    /// Encodes a dense [`Chunk`] into its palette-compressed form.
    #[must_use]
    pub fn from_chunk(chunk: &Chunk) -> Self {
        // Assign a palette slot to each distinct material in first-encounter order.
        let mut lookup: HashMap<BlockId, usize> = HashMap::new();
        let mut palette: Vec<BlockId> = Vec::new();

        for &block in &*chunk.blocks {
            if let Entry::Vacant(entry) = lookup.entry(block) {
                entry.insert(palette.len());
                palette.push(block);
            }
        }

        let bits_per_index = Self::bits_for_palette(palette.len());
        let mut indices = vec![0u64; Self::packed_word_count(bits_per_index)];

        for (voxel, &block) in chunk.blocks.iter().enumerate() {
            // Every block was inserted into `lookup` above, so this cannot miss.
            let index = lookup[&block] as u64;
            Self::write_packed(&mut indices, voxel, bits_per_index, index);
        }

        Self {
            palette,
            indices,
            bits_per_index,
        }
    }

    /// Decodes the palette-compressed form back into a dense [`Chunk`].
    #[must_use]
    pub fn to_chunk(&self) -> Chunk {
        let mut blocks = vec![BlockId::AIR; CHUNK_VOLUME].into_boxed_slice();
        for (voxel, slot) in blocks.iter_mut().enumerate() {
            // A stored index was produced from a palette position, so it is always in range for `palette`.
            #[expect(
                clippy::cast_possible_truncation,
                reason = "a packed index originates from a valid palette position and fits usize"
            )]
            let index = Self::read_packed(&self.indices, voxel, self.bits_per_index) as usize;
            *slot = self.palette[index];
        }
        Chunk { blocks }
    }

    /// Returns the number of bits required to store a palette index for a palette of `len` entries.
    fn bits_for_palette(len: usize) -> u32 {
        if len <= 1 {
            1
        } else {
            // Bit width of the largest index `len - 1`.
            usize::BITS - (len - 1).leading_zeros()
        }
    }

    /// Returns the number of 64-bit words needed to pack every voxel index at `bits_per_index` bits.
    fn packed_word_count(bits_per_index: u32) -> usize {
        let total_bits = CHUNK_VOLUME * bits_per_index as usize;
        total_bits.div_ceil(u64::BITS as usize)
    }

    /// Writes `value` as the `voxel`-th index of width `bits_per_index` into the packed `words` buffer.
    fn write_packed(words: &mut [u64], voxel: usize, bits_per_index: u32, value: u64) {
        let bit_offset = voxel * bits_per_index as usize;
        let word = bit_offset / 64;
        let bit_in_word = bit_offset % 64;

        // Low part: the bits that fit in the current word at `bit_in_word` and above. Bits shifted past bit 63 are dropped and handled by the spill below.
        words[word] |= value << bit_in_word;

        // High part: only present when the index straddles the word boundary.
        let bits_in_first = 64 - bit_in_word;
        if bits_per_index as usize > bits_in_first {
            words[word + 1] |= value >> bits_in_first;
        }
    }

    /// Reads the `voxel`-th index of width `bits_per_index` from the packed `words` buffer, reassembling any value that straddles a 64-bit word boundary.
    fn read_packed(words: &[u64], voxel: usize, bits_per_index: u32) -> u64 {
        let bit_offset = voxel * bits_per_index as usize;
        let word = bit_offset / 64;
        let bit_in_word = bit_offset % 64;

        // `bits_per_index` never reaches 64 (a chunk holds at most `CHUNK_VOLUME` distinct materials, so at most 15 bits), so this shift cannot overflow.
        let mask = (1u64 << bits_per_index) - 1;
        let bits_in_first = 64 - bit_in_word;

        let low = words[word] >> bit_in_word;
        if bits_per_index as usize <= bits_in_first {
            low & mask
        } else {
            // Reassemble a straddling index: low bits from the current word, high bits from the next.
            let high = words[word + 1] << bits_in_first;
            (low | high) & mask
        }
    }
}

#[cfg(test)]
#[path = "../tests/chunk.rs"]
mod tests;