//! Chunk management for distributed storage
//!
//! This module handles splitting large objects into fixed-size chunks
//! and reassembling them back into the original data.

use crate::config::ChunkConfig;

const TARGET_CHUNK_COUNT_PER_OBJECT: usize = 8;

/// Manages chunk operations for large objects
#[derive(Debug, Clone)]
pub struct ChunkManager {
    config: ChunkConfig,
}

impl ChunkManager {
    /// Create a new chunk manager with the given configuration
    pub fn new(config: ChunkConfig) -> Self {
        Self { config }
    }

    /// Create a new chunk manager with default configuration
    pub fn with_defaults() -> Self {
        Self::new(ChunkConfig::default())
    }

    /// Get the chunk size in bytes
    pub fn chunk_size(&self) -> usize {
        self.config.chunk_size
    }

    /// Choose the effective chunk size for an object of the given size.
    ///
    /// Small objects keep the configured default chunk size. Larger objects
    /// scale up to keep per-object chunk counts bounded without exceeding the
    /// configured maximum.
    pub fn effective_chunk_size(&self, total_size: usize) -> usize {
        if total_size == 0 {
            return self.config.chunk_size;
        }

        let min_chunk_size = self.config.min_chunk_size.min(self.config.chunk_size).max(1);
        let max_chunk_size = self.config.max_chunk_size.max(self.config.chunk_size);
        let required = total_size.div_ceil(TARGET_CHUNK_COUNT_PER_OBJECT);
        let alignment = min_chunk_size;
        let aligned_required = required.div_ceil(alignment) * alignment;

        aligned_required
            .max(self.config.chunk_size)
            .clamp(min_chunk_size, max_chunk_size)
    }

    /// Split data into chunks
    ///
    /// Returns a vector of chunks. Each chunk is at most `chunk_size` bytes,
    /// except the last chunk which may be smaller.
    pub fn split(&self, data: &[u8]) -> Vec<Vec<u8>> {
        self.split_with_chunk_size(data, self.config.chunk_size)
    }

    /// Split data into chunks using an explicit chunk size.
    pub fn split_with_chunk_size(&self, data: &[u8], chunk_size: usize) -> Vec<Vec<u8>> {
        if data.is_empty() {
            return vec![vec![]];
        }

        data.chunks(chunk_size).map(|c| c.to_vec()).collect()
    }

    /// Reassemble chunks into original data
    ///
    /// Chunks must be in order and complete.
    pub fn reassemble(&self, chunks: Vec<Vec<u8>>) -> Vec<u8> {
        chunks.into_iter().flatten().collect()
    }

    /// Calculate the number of chunks for a given data size
    pub fn chunk_count(&self, size: usize) -> usize {
        Self::chunk_count_for_size(size, self.config.chunk_size)
    }

    pub fn chunk_count_for_size(size: usize, chunk_size: usize) -> usize {
        if size == 0 {
            return 1;
        }
        size.div_ceil(chunk_size)
    }

    /// Calculate the size of a specific chunk
    ///
    /// Returns the size of the chunk at the given index for data of the given total size.
    pub fn chunk_size_at(&self, total_size: usize, chunk_index: usize) -> usize {
        Self::chunk_size_at_for_size(total_size, chunk_index, self.config.chunk_size)
    }

    pub fn chunk_size_at_for_size(
        total_size: usize,
        chunk_index: usize,
        chunk_size: usize,
    ) -> usize {
        let full_chunks = total_size / chunk_size;
        let remainder = total_size % chunk_size;

        if chunk_index < full_chunks {
            chunk_size
        } else if chunk_index == full_chunks && remainder > 0 {
            remainder
        } else {
            0
        }
    }

    /// Calculate the byte range for a specific chunk
    ///
    /// Returns (start_offset, length) for the chunk at the given index.
    pub fn chunk_range(&self, total_size: usize, chunk_index: usize) -> (usize, usize) {
        Self::chunk_range_for_size(total_size, chunk_index, self.config.chunk_size)
    }

    pub fn chunk_range_for_size(
        total_size: usize,
        chunk_index: usize,
        chunk_size: usize,
    ) -> (usize, usize) {
        let start = chunk_index * chunk_size;
        let length = Self::chunk_size_at_for_size(total_size, chunk_index, chunk_size);
        (start, length)
    }
}

impl Default for ChunkManager {
    fn default() -> Self {
        Self::with_defaults()
    }
}

/// Represents a chunk identifier
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct ChunkId {
    /// The object ID this chunk belongs to
    pub object_id: String,
    /// The chunk index within the object
    pub chunk_index: usize,
    /// The shard index (for erasure coding)
    pub shard_index: usize,
    /// Whether this is a parity shard
    pub is_parity: bool,
}

impl ChunkId {
    /// Create a new chunk ID for a simple (non-sharded) chunk
    pub fn simple(object_id: impl Into<String>, chunk_index: usize) -> Self {
        Self {
            object_id: object_id.into(),
            chunk_index,
            shard_index: 0,
            is_parity: false,
        }
    }

    /// Create a new chunk ID for a data shard
    pub fn data_shard(
        object_id: impl Into<String>,
        chunk_index: usize,
        shard_index: usize,
    ) -> Self {
        Self {
            object_id: object_id.into(),
            chunk_index,
            shard_index,
            is_parity: false,
        }
    }

    /// Create a new chunk ID for a parity shard
    pub fn parity_shard(
        object_id: impl Into<String>,
        chunk_index: usize,
        shard_index: usize,
    ) -> Self {
        Self {
            object_id: object_id.into(),
            chunk_index,
            shard_index,
            is_parity: true,
        }
    }

    /// Convert to a string key for storage
    pub fn to_key(&self) -> String {
        format!(
            "{}_{}_{}_{}",
            self.object_id,
            self.chunk_index,
            self.shard_index,
            if self.is_parity { "p" } else { "d" }
        )
    }

    /// Parse from a string key
    pub fn from_key(key: &str) -> Option<Self> {
        let parts: Vec<&str> = key.rsplitn(4, '_').collect();
        if parts.len() != 4 {
            return None;
        }

        let is_parity = parts[0] == "p";
        let shard_index = parts[1].parse().ok()?;
        let chunk_index = parts[2].parse().ok()?;
        let object_id = parts[3].to_string();

        Some(Self {
            object_id,
            chunk_index,
            shard_index,
            is_parity,
        })
    }
}

impl std::fmt::Display for ChunkId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.to_key())
    }
}

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

    #[test]
    fn test_split_empty() {
        let manager = ChunkManager::with_defaults();
        let chunks = manager.split(&[]);
        assert_eq!(chunks.len(), 1);
        assert!(chunks[0].is_empty());
    }

    #[test]
    fn test_split_smaller_than_chunk() {
        let manager = ChunkManager::new(ChunkConfig::new(1024));
        let data = vec![42u8; 512];
        let chunks = manager.split(&data);
        assert_eq!(chunks.len(), 1);
        assert_eq!(chunks[0], data);
    }

    #[test]
    fn test_split_exact_chunk_boundary() {
        let manager = ChunkManager::new(ChunkConfig::new(1024));
        let data = vec![42u8; 2048];
        let chunks = manager.split(&data);
        assert_eq!(chunks.len(), 2);
        assert_eq!(chunks[0].len(), 1024);
        assert_eq!(chunks[1].len(), 1024);
    }

    #[test]
    fn test_split_partial_last_chunk() {
        let manager = ChunkManager::new(ChunkConfig::new(1024));
        let data = vec![42u8; 1500];
        let chunks = manager.split(&data);
        assert_eq!(chunks.len(), 2);
        assert_eq!(chunks[0].len(), 1024);
        assert_eq!(chunks[1].len(), 476);
    }

    #[test]
    fn test_reassemble_preserves_data() {
        let manager = ChunkManager::new(ChunkConfig::new(1024));
        let original: Vec<u8> = (0..2500).map(|i| (i % 256) as u8).collect();
        let chunks = manager.split(&original);
        let reassembled = manager.reassemble(chunks);
        assert_eq!(original, reassembled);
    }

    #[test]
    fn test_chunk_count() {
        let manager = ChunkManager::new(ChunkConfig::new(1024));
        assert_eq!(manager.chunk_count(0), 1);
        assert_eq!(manager.chunk_count(512), 1);
        assert_eq!(manager.chunk_count(1024), 1);
        assert_eq!(manager.chunk_count(1025), 2);
        assert_eq!(manager.chunk_count(2048), 2);
        assert_eq!(manager.chunk_count(3000), 3);
    }

    #[test]
    fn test_chunk_size_at() {
        let manager = ChunkManager::new(ChunkConfig::new(1024));
        // 2500 bytes = 2 full chunks (1024) + 1 partial (452)
        assert_eq!(manager.chunk_size_at(2500, 0), 1024);
        assert_eq!(manager.chunk_size_at(2500, 1), 1024);
        assert_eq!(manager.chunk_size_at(2500, 2), 452);
        assert_eq!(manager.chunk_size_at(2500, 3), 0);
    }

    #[test]
    fn test_chunk_range() {
        let manager = ChunkManager::new(ChunkConfig::new(1024));
        assert_eq!(manager.chunk_range(2500, 0), (0, 1024));
        assert_eq!(manager.chunk_range(2500, 1), (1024, 1024));
        assert_eq!(manager.chunk_range(2500, 2), (2048, 452));
    }

    #[test]
    fn test_effective_chunk_size_scales_large_objects_up_to_target_chunk_count() {
        let manager = ChunkManager::default();

        assert_eq!(manager.effective_chunk_size(4 * 1024 * 1024), 8 * 1024 * 1024);
        assert_eq!(manager.effective_chunk_size(256 * 1024 * 1024), 32 * 1024 * 1024);
        assert_eq!(manager.effective_chunk_size(1024 * 1024 * 1024), 64 * 1024 * 1024);
    }

    #[test]
    fn test_chunk_id_to_key() {
        let id = ChunkId::data_shard("obj123", 0, 2);
        assert_eq!(id.to_key(), "obj123_0_2_d");

        let id = ChunkId::parity_shard("obj123", 1, 4);
        assert_eq!(id.to_key(), "obj123_1_4_p");
    }

    #[test]
    fn test_chunk_id_roundtrip() {
        let original = ChunkId::data_shard("my-object", 5, 3);
        let key = original.to_key();
        let parsed = ChunkId::from_key(&key).unwrap();
        assert_eq!(original, parsed);
    }
}
