use crate::error::{Error, Result};
use crate::filters;
use crate::format::fixed_array::FixedArrayElement;
use crate::io::reader::HdfReader;
use std::io::{Read, Seek};
use std::thread;

use super::chunk_read::ChunkReadContext;
use super::{u64_from_usize, usize_from_u64, Dataset, DatasetInfo};

const MIN_PARALLEL_LINEAR_DEFLATE_CHUNKS_1D: usize = 8;
const MIN_PARALLEL_LINEAR_DEFLATE_BYTES_1D: usize = 64 * 1024;

impl Dataset {
    pub(super) fn read_chunked_fixed_array<R: Read + Seek>(
        reader: &mut HdfReader<R>,
        info: &DatasetInfo,
        chunk_ctx: &ChunkReadContext<'_>,
    ) -> Result<Vec<u8>> {
        let mut output = Self::scratch_output(chunk_ctx.total_bytes);
        Self::read_chunked_fixed_array_into(reader, info, chunk_ctx, &mut output)?;
        Ok(output)
    }

    pub(super) fn read_chunked_fixed_array_into<R: Read + Seek>(
        reader: &mut HdfReader<R>,
        info: &DatasetInfo,
        chunk_ctx: &ChunkReadContext<'_>,
        output: &mut [u8],
    ) -> Result<()> {
        if output.len() != chunk_ctx.total_bytes {
            return Err(Error::InvalidFormat(format!(
                "fixed-array chunk output buffer has {} bytes, expected {}",
                output.len(),
                chunk_ctx.total_bytes
            )));
        }
        let filtered = info
            .filter_pipeline
            .as_ref()
            .map(|pipeline| !pipeline.filters.is_empty())
            .unwrap_or(false);
        let chunk_size_len = if filtered {
            Self::filtered_chunk_size_len(
                info,
                chunk_ctx.chunk_bytes,
                usize::from(reader.sizeof_size()),
            )?
        } else {
            0
        };

        let mut elements = Vec::new();
        crate::format::fixed_array::read_fixed_array_chunks_into(
            reader,
            chunk_ctx.idx_addr,
            filtered,
            chunk_size_len,
            &mut elements,
        )?;
        let chunks_per_dim = Self::chunks_per_dim(chunk_ctx.data_dims, chunk_ctx.chunk_dims)?;
        if !Self::has_full_linear_chunk_coverage_1d(elements.iter(), chunk_ctx)? {
            Self::filled_data_into(
                chunk_ctx.total_bytes / chunk_ctx.element_size,
                chunk_ctx.element_size,
                info,
                output,
            )?;
        }
        let mut compressed_scratch = Vec::new();
        let mut shuffle_scratch = Vec::new();
        let mut raw_scratch = Vec::new();
        let mut coords = Vec::with_capacity(chunk_ctx.data_dims.len());
        let mut handled = if Self::linear_index_uses_unfiltered_coalescing(info) {
            Self::try_read_coalesced_linear_index_unfiltered_chunks_1d(
                reader,
                info,
                chunk_ctx,
                elements.iter(),
                "fixed-array chunk size",
                output,
            )?
        } else {
            Vec::new()
        };
        let logical_chunks = Self::expected_chunk_count_1d(chunk_ctx)?.unwrap_or(elements.len());
        if handled.is_empty() {
            handled.resize(logical_chunks.min(elements.len()), false);
        }
        Self::try_read_parallel_deflate_linear_chunks_1d(
            reader,
            info,
            chunk_ctx,
            &elements,
            output,
            &mut handled,
            "fixed-array chunk size",
        )?;

        for (chunk_index, element) in elements.iter().take(logical_chunks).enumerate() {
            Self::trace_linear_chunk_lookup(
                "hdf5.chunk_index.fixed_array.lookup",
                chunk_ctx.idx_addr,
                u64_from_usize(chunk_index, "fixed-array chunk index")?,
                element.addr,
                element.nbytes.unwrap_or(u64_from_usize(
                    chunk_ctx.chunk_bytes,
                    "fixed-array chunk size",
                )?),
                element.filter_mask,
            );

            if crate::io::reader::is_undef_addr(element.addr) {
                continue;
            }

            if handled.get(chunk_index).copied().unwrap_or(false) {
                continue;
            }
            Self::implicit_chunk_coords_into(
                chunk_index,
                chunk_ctx.chunk_dims,
                &chunks_per_dim,
                &mut coords,
            )?;
            let read_size = usize_from_u64(
                element.nbytes.unwrap_or(u64_from_usize(
                    chunk_ctx.chunk_bytes,
                    "fixed-array chunk size",
                )?),
                "fixed-array chunk size",
            )?;
            if Self::try_read_full_chunk_1d_into_output(
                reader,
                info,
                chunk_ctx,
                &coords,
                element.addr,
                read_size,
                element.filter_mask,
                output,
                &mut compressed_scratch,
                &mut shuffle_scratch,
            )? {
                continue;
            }
            Self::read_linear_chunk_payload_into_scratch(
                reader,
                element.addr,
                read_size,
                "fixed-array",
                chunk_index,
                &mut raw_scratch,
            )?;

            if let Some(ref pipeline) = info.filter_pipeline {
                if !pipeline.filters.is_empty() {
                    let filtered = filters::apply_pipeline_reverse_with_mask_expected(
                        &raw_scratch,
                        pipeline,
                        chunk_ctx.element_size,
                        element.filter_mask,
                        chunk_ctx.chunk_bytes,
                    )?;
                    Self::copy_chunk_to_output(
                        &filtered,
                        &coords,
                        chunk_ctx.data_dims,
                        chunk_ctx.chunk_dims,
                        chunk_ctx.element_size,
                        output,
                    )?;
                    continue;
                }
            }

            Self::copy_chunk_to_output(
                &raw_scratch,
                &coords,
                chunk_ctx.data_dims,
                chunk_ctx.chunk_dims,
                chunk_ctx.element_size,
                output,
            )?;
        }

        Ok(())
    }

    pub(super) fn read_chunked_extensible_array<R: Read + Seek>(
        reader: &mut HdfReader<R>,
        info: &DatasetInfo,
        chunk_ctx: &ChunkReadContext<'_>,
    ) -> Result<Vec<u8>> {
        let mut output = Self::scratch_output(chunk_ctx.total_bytes);
        Self::read_chunked_extensible_array_into(reader, info, chunk_ctx, &mut output)?;
        Ok(output)
    }

    pub(super) fn read_chunked_extensible_array_into<R: Read + Seek>(
        reader: &mut HdfReader<R>,
        info: &DatasetInfo,
        chunk_ctx: &ChunkReadContext<'_>,
        output: &mut [u8],
    ) -> Result<()> {
        if output.len() != chunk_ctx.total_bytes {
            return Err(Error::InvalidFormat(format!(
                "extensible-array chunk output buffer has {} bytes, expected {}",
                output.len(),
                chunk_ctx.total_bytes
            )));
        }
        let filtered = info
            .filter_pipeline
            .as_ref()
            .map(|pipeline| !pipeline.filters.is_empty())
            .unwrap_or(false);
        let chunk_size_len = if filtered {
            Self::filtered_chunk_size_len(
                info,
                chunk_ctx.chunk_bytes,
                usize::from(reader.sizeof_size()),
            )?
        } else {
            0
        };

        let mut elements = Vec::new();
        crate::format::extensible_array::read_extensible_array_chunks_into(
            reader,
            chunk_ctx.idx_addr,
            filtered,
            chunk_size_len,
            &mut elements,
        )?;
        let chunks_per_dim = Self::chunks_per_dim(chunk_ctx.data_dims, chunk_ctx.chunk_dims)?;
        if !Self::has_full_linear_chunk_coverage_1d(elements.iter(), chunk_ctx)? {
            Self::filled_data_into(
                chunk_ctx.total_bytes / chunk_ctx.element_size,
                chunk_ctx.element_size,
                info,
                output,
            )?;
        }
        let mut compressed_scratch = Vec::new();
        let mut shuffle_scratch = Vec::new();
        let mut raw_scratch = Vec::new();
        let mut coords = Vec::with_capacity(chunk_ctx.data_dims.len());
        let mut handled = if Self::linear_index_uses_unfiltered_coalescing(info) {
            Self::try_read_coalesced_linear_index_unfiltered_chunks_1d(
                reader,
                info,
                chunk_ctx,
                elements.iter(),
                "extensible-array chunk size",
                output,
            )?
        } else {
            Vec::new()
        };
        let logical_chunks = Self::expected_chunk_count_1d(chunk_ctx)?.unwrap_or(elements.len());
        if handled.is_empty() {
            handled.resize(logical_chunks.min(elements.len()), false);
        }
        Self::try_read_parallel_deflate_linear_chunks_1d(
            reader,
            info,
            chunk_ctx,
            &elements,
            output,
            &mut handled,
            "extensible-array chunk size",
        )?;

        for (chunk_index, element) in elements.iter().take(logical_chunks).enumerate() {
            Self::trace_linear_chunk_lookup(
                "hdf5.chunk_index.extensible_array.lookup",
                chunk_ctx.idx_addr,
                u64_from_usize(chunk_index, "extensible-array chunk index")?,
                element.addr,
                element.nbytes.unwrap_or(u64_from_usize(
                    chunk_ctx.chunk_bytes,
                    "extensible-array chunk size",
                )?),
                element.filter_mask,
            );

            if crate::io::reader::is_undef_addr(element.addr) {
                continue;
            }

            if handled.get(chunk_index).copied().unwrap_or(false) {
                continue;
            }
            Self::implicit_chunk_coords_into(
                chunk_index,
                chunk_ctx.chunk_dims,
                &chunks_per_dim,
                &mut coords,
            )?;
            let read_size = usize_from_u64(
                element.nbytes.unwrap_or(u64_from_usize(
                    chunk_ctx.chunk_bytes,
                    "extensible-array chunk size",
                )?),
                "extensible-array chunk size",
            )?;
            if Self::try_read_full_chunk_1d_into_output(
                reader,
                info,
                chunk_ctx,
                &coords,
                element.addr,
                read_size,
                element.filter_mask,
                output,
                &mut compressed_scratch,
                &mut shuffle_scratch,
            )? {
                continue;
            }
            Self::read_linear_chunk_payload_into_scratch(
                reader,
                element.addr,
                read_size,
                "extensible-array",
                chunk_index,
                &mut raw_scratch,
            )?;

            if let Some(ref pipeline) = info.filter_pipeline {
                if !pipeline.filters.is_empty() {
                    let filtered = filters::apply_pipeline_reverse_with_mask_expected(
                        &raw_scratch,
                        pipeline,
                        chunk_ctx.element_size,
                        element.filter_mask,
                        chunk_ctx.chunk_bytes,
                    )?;
                    Self::copy_chunk_to_output(
                        &filtered,
                        &coords,
                        chunk_ctx.data_dims,
                        chunk_ctx.chunk_dims,
                        chunk_ctx.element_size,
                        output,
                    )?;
                    continue;
                }
            }

            Self::copy_chunk_to_output(
                &raw_scratch,
                &coords,
                chunk_ctx.data_dims,
                chunk_ctx.chunk_dims,
                chunk_ctx.element_size,
                output,
            )?;
        }

        Ok(())
    }

    pub(super) fn has_full_linear_chunk_coverage_1d<'a, I>(
        elements: I,
        chunk_ctx: &ChunkReadContext<'_>,
    ) -> Result<bool>
    where
        I: IntoIterator<Item = &'a crate::format::fixed_array::FixedArrayElement>,
    {
        let Some(expected_chunks) = Self::expected_chunk_count_1d(chunk_ctx)? else {
            return Ok(false);
        };
        let mut count = 0usize;
        for element in elements {
            if crate::io::reader::is_undef_addr(element.addr) {
                return Ok(false);
            }
            count = count
                .checked_add(1)
                .ok_or_else(|| Error::InvalidFormat("chunk count overflow".into()))?;
        }
        Ok(count == expected_chunks)
    }

    fn linear_index_uses_unfiltered_coalescing(info: &DatasetInfo) -> bool {
        info.filter_pipeline
            .as_ref()
            .map(|pipeline| pipeline.filters.is_empty())
            .unwrap_or(true)
    }

    fn try_read_parallel_deflate_linear_chunks_1d<R: Read + Seek>(
        reader: &mut HdfReader<R>,
        info: &DatasetInfo,
        chunk_ctx: &ChunkReadContext<'_>,
        elements: &[FixedArrayElement],
        output: &mut [u8],
        handled: &mut [bool],
        size_context: &'static str,
    ) -> Result<()> {
        // libhdf5 reads each filtered chunk through the raw data chunk cache
        // and runs H5Z_pipeline serially in-process. This Rust-only fast path
        // keeps that serial behavior as the semantic baseline but decompresses
        // independent deflate-only 1D full-prefix chunks on scoped workers.
        let Some(pipeline) = info.filter_pipeline.as_ref() else {
            return Ok(());
        };
        if !Self::is_deflate_only_pipeline(pipeline) || chunk_ctx.data_dims.len() != 1 {
            return Ok(());
        }
        if chunk_ctx.chunk_dims.len() != 1 || output.len() != chunk_ctx.total_bytes {
            return Ok(());
        }

        let data_size = usize_from_u64(chunk_ctx.data_dims[0], "dataset dimension")?;
        let chunk_size = usize_from_u64(chunk_ctx.chunk_dims[0], "chunk dimension")?;
        if chunk_size == 0 {
            return Ok(());
        }
        let full_chunk_count = data_size / chunk_size;
        if full_chunk_count < MIN_PARALLEL_LINEAR_DEFLATE_CHUNKS_1D
            || elements.len() < full_chunk_count
        {
            return Ok(());
        }
        let chunk_bytes = chunk_size
            .checked_mul(chunk_ctx.element_size)
            .ok_or_else(|| Error::InvalidFormat("chunk byte size overflow".into()))?;
        if chunk_bytes == 0 || chunk_bytes != chunk_ctx.chunk_bytes {
            return Ok(());
        }
        let full_output_len = full_chunk_count
            .checked_mul(chunk_bytes)
            .ok_or_else(|| Error::InvalidFormat("parallel chunk output length overflow".into()))?;
        if full_output_len > output.len() || full_output_len < MIN_PARALLEL_LINEAR_DEFLATE_BYTES_1D
        {
            return Ok(());
        }

        let worker_count = super::support::parallel_deflate_worker_count(full_chunk_count);
        if worker_count <= 1 {
            return Ok(());
        }

        let mut payloads = Vec::with_capacity(full_chunk_count);
        for (chunk_index, element) in elements.iter().take(full_chunk_count).enumerate() {
            if handled.get(chunk_index).copied().unwrap_or(false)
                || element.filter_mask != 0
                || crate::io::reader::is_undef_addr(element.addr)
            {
                return Ok(());
            }
            let read_size = usize_from_u64(
                element
                    .nbytes
                    .unwrap_or(u64_from_usize(chunk_bytes, size_context)?),
                size_context,
            )?;
            let mut payload = vec![0u8; read_size];
            reader.seek(element.addr)?;
            reader.read_bytes_into(&mut payload).map_err(|err| {
                Error::InvalidFormat(format!(
                    "failed to read {size_context} chunk {chunk_index} at address {} with size {read_size}: {err}",
                    element.addr
                ))
            })?;
            payloads.push(payload);
        }

        let chunks_per_worker = full_chunk_count.div_ceil(worker_count);
        let full_output = &mut output[..full_output_len];
        let parallel_result: Result<()> = thread::scope(|scope| {
            let mut handles = Vec::new();
            let mut remaining_output = full_output;
            let mut payload_start = 0usize;
            while payload_start < full_chunk_count {
                let payload_end = (payload_start + chunks_per_worker).min(full_chunk_count);
                let group_chunks = payload_end - payload_start;
                let group_bytes = group_chunks * chunk_bytes;
                let (group_output, next_output) = remaining_output.split_at_mut(group_bytes);
                let group_payloads = &payloads[payload_start..payload_end];
                handles.push(scope.spawn(move || -> Result<()> {
                    for (chunk_offset, payload) in group_payloads.iter().enumerate() {
                        let start = chunk_offset * chunk_bytes;
                        let end = start + chunk_bytes;
                        crate::filters::deflate::decompress_exact_into(
                            payload,
                            &mut group_output[start..end],
                        )?;
                    }
                    Ok(())
                }));
                remaining_output = next_output;
                payload_start = payload_end;
            }
            for handle in handles {
                handle.join().map_err(|_| {
                    Error::InvalidFormat("parallel deflate worker panicked".into())
                })??;
            }
            Ok(())
        });
        parallel_result?;
        super::support::record_parallel_deflate_chunks_handled(full_chunk_count);

        for flag in handled.iter_mut().take(full_chunk_count) {
            *flag = true;
        }
        Ok(())
    }

    fn read_linear_chunk_payload_into_scratch<R: Read + Seek>(
        reader: &mut HdfReader<R>,
        addr: u64,
        read_size: usize,
        index_name: &'static str,
        chunk_index: usize,
        scratch: &mut Vec<u8>,
    ) -> Result<()> {
        reader.seek(addr)?;
        scratch.resize(read_size, 0);
        reader.read_bytes_into(scratch).map_err(|err| {
            Error::InvalidFormat(format!(
                "failed to read {index_name} chunk {chunk_index} at address {addr} with size {read_size}: {err}"
            ))
        })?;
        Ok(())
    }

    #[cfg(feature = "tracehash")]
    fn trace_linear_chunk_lookup(
        function: &'static str,
        index_addr: u64,
        chunk_index: u64,
        addr: u64,
        nbytes: u64,
        filter_mask: u32,
    ) {
        let mut th = tracehash::Call::new(function, file!(), line!());
        th.input_u64(index_addr);
        th.input_u64(chunk_index);
        th.output_value(&(true));
        th.output_u64(addr);
        th.output_u64(if crate::io::reader::is_undef_addr(addr) {
            0
        } else {
            nbytes
        });
        th.output_u64(u64::from(filter_mask));
        th.finish();
    }

    #[cfg(not(feature = "tracehash"))]
    fn trace_linear_chunk_lookup(
        _function: &'static str,
        _index_addr: u64,
        _chunk_index: u64,
        _addr: u64,
        _nbytes: u64,
        _filter_mask: u32,
    ) {
    }
}
