use std::borrow::Cow;
use std::io::{Seek, SeekFrom, Write};

use crate::error::{Error, Result};
use crate::format::messages::data_layout::{ChunkIndexType, LayoutClass};
use crate::format::messages::filter_pipeline::{FILTER_DEFLATE, FILTER_FLETCHER32, FILTER_SHUFFLE};

use super::MutableFile;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum WritableChunkIndexKind {
    BTreeV1,
    FixedArray,
    ExtensibleArray,
    BTreeV2,
}

impl MutableFile {
    /// Write a full uncompressed chunk and update the dataset's chunk index.
    ///
    /// This supports chunked datasets written by this crate where the chunk
    /// index variant is one of the currently mutable writer-side subsets.
    pub fn write_chunk(
        &mut self,
        dataset_path: &str,
        chunk_coords: &[u64],
        data: &[u8],
    ) -> Result<()> {
        let ds = self.dataset(dataset_path)?;
        let info = ds.info()?;

        if info.layout.layout_class != LayoutClass::Chunked {
            return Err(Error::InvalidFormat(
                "write_chunk only supports chunked datasets".into(),
            ));
        }
        if info.layout.version > 3
            && !matches!(
                info.layout.chunk_index_type,
                Some(ChunkIndexType::BTreeV1)
                    | Some(ChunkIndexType::FixedArray)
                    | Some(ChunkIndexType::ExtensibleArray)
                    | Some(ChunkIndexType::BTreeV2)
            )
        {
            return Err(Error::Unsupported(
                "write_chunk currently supports only v1 B-tree, fixed-array, simple extensible-array, and simple v2 B-tree chunk indexes".into(),
            ));
        }

        let chunk_data_dims = Self::chunk_data_dims(&info)?;
        Self::validate_chunk_coords(chunk_coords, chunk_data_dims)?;

        let element_size = Self::u64_to_usize(u64::from(info.datatype.size), "datatype size")?;
        let expected_len = Self::expected_chunk_len(chunk_data_dims, element_size)?;
        Self::validate_chunk_write_len(data.len(), expected_len)?;
        let filtered = Self::encode_chunk_write_data(&info, data, element_size)?;

        let index_addr = info
            .layout
            .chunk_index_addr
            .ok_or_else(|| Error::InvalidFormat("chunked dataset missing B-tree address".into()))?;

        let chunk_addr = self.write_handle.seek(SeekFrom::End(0))?;
        self.write_handle.write_all(&filtered)?;
        self.rewrite_chunk_index(
            index_addr,
            &info,
            chunk_coords,
            chunk_data_dims,
            Self::usize_to_u64(filtered.len(), "filtered chunk size")?,
            chunk_addr,
            expected_len,
            element_size,
        )?;
        self.write_handle.flush()?;
        self.reopen_reader()?;

        Ok(())
    }

    fn chunk_data_dims(info: &crate::hl::dataset::DatasetInfo) -> Result<&[u64]> {
        let chunk_dims = info
            .layout
            .chunk_dims
            .as_ref()
            .ok_or_else(|| Error::InvalidFormat("chunked layout missing chunk dims".into()))?;
        if chunk_dims.len() == info.dataspace.dims.len() + 1 {
            Ok(&chunk_dims[..info.dataspace.dims.len()])
        } else if chunk_dims.len() == info.dataspace.dims.len() {
            Ok(chunk_dims)
        } else {
            Err(Error::InvalidFormat(format!(
                "chunk dimension rank {} does not match dataset rank {}",
                chunk_dims.len(),
                info.dataspace.dims.len()
            )))
        }
    }

    fn validate_chunk_coords(chunk_coords: &[u64], chunk_data_dims: &[u64]) -> Result<()> {
        if chunk_coords.len() != chunk_data_dims.len() {
            return Err(Error::InvalidFormat(format!(
                "chunk coordinate rank {} does not match dataset rank {}",
                chunk_coords.len(),
                chunk_data_dims.len()
            )));
        }
        for ((idx, &coord), &chunk) in chunk_coords.iter().enumerate().zip(chunk_data_dims) {
            if chunk == 0 || coord % chunk != 0 {
                return Err(Error::InvalidFormat(format!(
                    "chunk coordinate {idx}={coord} is not aligned to chunk size {chunk}"
                )));
            }
        }
        Ok(())
    }

    fn expected_chunk_len(chunk_data_dims: &[u64], element_size: usize) -> Result<usize> {
        let chunk_elements = chunk_data_dims.iter().try_fold(1usize, |acc, &dim| {
            let dim = Self::u64_to_usize(dim, "chunk dimension")?;
            acc.checked_mul(dim)
                .ok_or_else(|| Error::InvalidFormat("chunk element count overflow".into()))
        })?;
        chunk_elements
            .checked_mul(element_size)
            .ok_or_else(|| Error::InvalidFormat("chunk byte size overflow".into()))
    }

    fn validate_chunk_write_len(actual_len: usize, expected_len: usize) -> Result<()> {
        if actual_len != expected_len {
            return Err(Error::InvalidFormat(format!(
                "chunk data has {actual_len} bytes, expected {expected_len}",
            )));
        }
        Ok(())
    }

    fn encode_chunk_write_data<'a>(
        info: &crate::hl::dataset::DatasetInfo,
        data: &'a [u8],
        element_size: usize,
    ) -> Result<Cow<'a, [u8]>> {
        let mut filtered = Cow::Borrowed(data);
        let mut filter_buf = Vec::new();
        if let Some(ref pipeline) = info.filter_pipeline {
            for filter in &pipeline.filters {
                match filter.id {
                    FILTER_SHUFFLE => {
                        filter_buf.clear();
                        filter_buf.resize(filtered.len(), 0);
                        crate::filters::shuffle::shuffle_into(
                            filtered.as_ref(),
                            element_size,
                            &mut filter_buf,
                        )?;
                        Self::replace_encoded_filter_buffer(&mut filtered, &mut filter_buf);
                    }
                    FILTER_DEFLATE => {
                        let level = filter.client_data.first().copied().unwrap_or(6);
                        filter_buf.clear();
                        crate::filters::deflate::compress_into(
                            filtered.as_ref(),
                            level,
                            &mut filter_buf,
                        )?;
                        Self::replace_encoded_filter_buffer(&mut filtered, &mut filter_buf);
                    }
                    FILTER_FLETCHER32 => {
                        crate::filters::fletcher32::append_checksum_in_place(filtered.to_mut())?;
                    }
                    other => {
                        return Err(Error::Unsupported(format!(
                            "write_chunk cannot encode filter {other}"
                        )));
                    }
                }
            }
        }
        Ok(filtered)
    }

    fn replace_encoded_filter_buffer<'a>(filtered: &mut Cow<'a, [u8]>, filter_buf: &mut Vec<u8>) {
        let previous = std::mem::replace(filtered, Cow::Owned(std::mem::take(filter_buf)));
        if let Cow::Owned(mut previous) = previous {
            previous.clear();
            *filter_buf = previous;
        }
    }

    fn rewrite_chunk_index(
        &mut self,
        index_addr: u64,
        info: &crate::hl::dataset::DatasetInfo,
        chunk_coords: &[u64],
        chunk_data_dims: &[u64],
        filtered_len: u64,
        chunk_addr: u64,
        expected_len: usize,
        element_size: usize,
    ) -> Result<()> {
        match Self::writable_chunk_index_kind(info) {
            WritableChunkIndexKind::FixedArray => self.rewrite_fixed_array_chunk(
                index_addr,
                &info,
                chunk_coords,
                chunk_data_dims,
                filtered_len,
                chunk_addr,
                expected_len,
            ),
            WritableChunkIndexKind::ExtensibleArray => self.rewrite_extensible_array_chunk(
                index_addr,
                &info,
                chunk_coords,
                chunk_data_dims,
                filtered_len,
                chunk_addr,
                expected_len,
            ),
            WritableChunkIndexKind::BTreeV2 => self.rewrite_btree_v2_chunk(
                index_addr,
                &info,
                chunk_coords,
                chunk_data_dims,
                filtered_len,
                chunk_addr,
                expected_len,
            ),
            WritableChunkIndexKind::BTreeV1 => self.rewrite_leaf_chunk_btree(
                index_addr,
                chunk_coords,
                Self::u64_to_u32(filtered_len, "filtered chunk size")?,
                chunk_addr,
                Self::usize_to_u32(element_size, "datatype size")?,
            ),
        }
    }

    fn writable_chunk_index_kind(info: &crate::hl::dataset::DatasetInfo) -> WritableChunkIndexKind {
        match info.layout.chunk_index_type {
            Some(ChunkIndexType::FixedArray) => WritableChunkIndexKind::FixedArray,
            Some(ChunkIndexType::ExtensibleArray) => WritableChunkIndexKind::ExtensibleArray,
            Some(ChunkIndexType::BTreeV2) => WritableChunkIndexKind::BTreeV2,
            _ => WritableChunkIndexKind::BTreeV1,
        }
    }
}

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

    #[test]
    fn expected_chunk_len_rejects_element_count_overflow() {
        let err = MutableFile::expected_chunk_len(&[u64::MAX, 2], 1).unwrap_err();
        assert!(
            err.to_string().contains("chunk element count"),
            "unexpected error: {err}"
        );
    }
}
