// (C) Copyright 2026- ECMWF and individual contributors.
//
// This software is licensed under the terms of the Apache Licence Version 2.0
// which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
// In applying this licence, ECMWF does not waive the privileges and immunities
// granted to it by virtue of its status as an intergovernmental organisation nor
// does it submit to any jurisdiction.

use crate::encode::build_pipeline_config_with_backend;
use crate::error::{Result, TensogramError};
use crate::framing;

use crate::types::{DataObjectDescriptor, DecodedObject, GlobalMetadata};
use tensogram_encodings::pipeline;

fn extract_block_offsets(
    params: &std::collections::BTreeMap<String, ciborium::Value>,
) -> Result<Vec<u64>> {
    match params.get("szip_block_offsets") {
        Some(ciborium::Value::Array(arr)) => arr
            .iter()
            .map(|v| match v {
                ciborium::Value::Integer(i) => {
                    let n: i128 = (*i).into();
                    u64::try_from(n).map_err(|_| {
                        TensogramError::Metadata("szip_block_offset out of u64 range".to_string())
                    })
                }
                _ => Err(TensogramError::Metadata(
                    "szip_block_offsets must contain integers".to_string(),
                )),
            })
            .collect(),
        Some(_) => Err(TensogramError::Metadata(
            "szip_block_offsets must be an array".to_string(),
        )),
        None => Err(TensogramError::Compression(
            "missing szip_block_offsets in payload metadata (required for partial range decode)"
                .to_string(),
        )),
    }
}

/// Options for decoding.
#[derive(Debug, Clone)]
pub struct DecodeOptions {
    /// When true (the default), decoded payloads are converted to the
    /// caller's native byte order regardless of the wire byte order declared
    /// in the descriptor.  Set to false to receive bytes in the message's
    /// declared wire byte order (rare — useful for zero-copy forwarding).
    pub native_byte_order: bool,
    /// Which backend to use for szip / zstd when both FFI and pure-Rust
    /// implementations are compiled in.
    pub compression_backend: pipeline::CompressionBackend,
    /// Thread budget for the multi-threaded decoding pipeline.
    ///
    /// Semantics match
    /// [`EncodeOptions.threads`](crate::encode::EncodeOptions::threads):
    /// `0` means sequential (may be overridden by `TENSOGRAM_THREADS`),
    /// `1` means explicit single-threaded execution, `N ≥ 2` builds a
    /// scoped pool.  Output bytes are byte-identical to the
    /// sequential path regardless of `N`.
    pub threads: u32,
    /// Minimum total payload bytes below which the parallel path is
    /// skipped.  See
    /// [`EncodeOptions.parallel_threshold_bytes`](crate::encode::EncodeOptions::parallel_threshold_bytes).
    pub parallel_threshold_bytes: Option<usize>,
    /// When `true` (the default) AND the object carries a
    /// `NTensorFrame` `masks` sub-map, decompress the masks
    /// and write the canonical NaN / +Inf / -Inf bit pattern at
    /// every `1` position in the decoded output.  See
    /// `plans/WIRE_FORMAT.md` §6.5.4 for the (lossy) reconstruction
    /// caveat — only the canonical quiet-NaN / ±∞ bit pattern is
    /// restored; specific NaN payloads are not preserved.
    ///
    /// Set to `false` to skip restoration and receive the
    /// `0.0`-substituted bytes as they are on disk.  Callers who
    /// need the raw masks alongside the substituted payload use
    /// [`decode_with_masks`] instead.
    pub restore_non_finite: bool,
    /// When `true`, the decoder verifies the inline xxh3-64 hash
    /// of every data-object frame it materialises against the
    /// recomputed digest of the frame body.  Verification is
    /// fused with decode — bytes are hashed while hot in
    /// cache/buffer, so the cost is one extra walk over the
    /// post-encoding payload.  See `plans/DESIGN.md`
    /// §"Integrity Hashing" and `plans/WIRE_FORMAT.md` §11.1.
    ///
    /// When `false` (the default), no verification is performed.
    /// Decode is a pure deserialisation; the per-frame
    /// `HASH_PRESENT` flag and the slot value are both ignored.
    ///
    /// Strict-input rules under `verify_hash = true`:
    ///   * frame's `HASH_PRESENT` flag clear → returns
    ///     [`TensogramError::MissingHash`] with the offending
    ///     object index.
    ///   * flag set + slot disagrees with recomputed digest →
    ///     returns [`TensogramError::HashMismatch`] with the
    ///     offending object index.
    ///
    /// Scope: applies to **data-object** frames only (i.e. every
    /// `NTensorFrame` the decoder touches).  Header / footer /
    /// index / hash / preceder frames are not verified by this
    /// option — that's the offline `tensogram validate
    /// --checksum` validator's responsibility.
    ///
    /// **Not respected by `decode_range` / `decode_range_from_frame`
    /// / `RemoteBackend::decode_range`.**  Range decode reads
    /// only a slice of the encoded payload; verifying the inline
    /// hash would require reading every byte the optimisation is
    /// designed to avoid.  When integrity matters on a partial
    /// read, call `decode_object` with `verify_hash=true` —
    /// that path materialises the full body anyway, so the
    /// verification is free.
    ///
    /// **Verify-first ordering.** [`decode`], [`decode_object`],
    /// and [`decode_object_from_frame`] all run hash verification
    /// *before* CBOR descriptor parsing.  A tamper anywhere in
    /// the hashed body region — encoded payload bytes, mask
    /// blobs, or the CBOR descriptor itself — surfaces uniformly
    /// as [`TensogramError::HashMismatch`] (or
    /// [`TensogramError::MissingHash`] when `HASH_PRESENT` is
    /// clear), never as a downstream
    /// [`TensogramError::Metadata`].  Catch-by-class semantics
    /// behave consistently across all three entry points.
    pub verify_hash: bool,
}

impl Default for DecodeOptions {
    fn default() -> Self {
        Self {
            native_byte_order: true,
            compression_backend: pipeline::CompressionBackend::default(),
            threads: 0,
            parallel_threshold_bytes: None,
            restore_non_finite: true,
            verify_hash: false,
        }
    }
}

/// Walk the data-object frames in a single message buffer and
/// verify each frame's inline xxh3 hash *before* any CBOR is
/// parsed downstream.
///
/// `target_index = None` verifies every data-object frame in the
/// message; `target_index = Some(i)` verifies only the frame at
/// the i-th data-object position (used by [`decode_object`] which
/// only materialises one object).  Either way, hash verification
/// runs before [`framing::decode_message`] gets to parse the CBOR
/// descriptor, so a tamper anywhere in the hashed body region —
/// including within the CBOR bytes — surfaces as
/// [`TensogramError::HashMismatch`] / [`TensogramError::MissingHash`]
/// rather than the [`TensogramError::Metadata`] that a downstream
/// CBOR parse would emit.
///
/// `buf` may begin with one or more concatenated `.tgm` messages
/// (the `multi_message.tgm` golden is the canonical example);
/// this helper verifies only the **first** message, matching the
/// scope of [`framing::decode_message`].  The walk runs from
/// `PREAMBLE_SIZE` up to `msg_end`, where `msg_end` is computed
/// from `Preamble.total_length` when non-zero (the closed-stream
/// case) and from `buf.len() - POSTAMBLE_SIZE` otherwise (the
/// open-stream case).  Without this bound, a `[hashed][unhashed]`
/// concatenation would have its second message's frames hashed by
/// the verify pre-pass, surfacing spurious `MissingHash` errors
/// even though `decode` only ever returns the first message's
/// contents.  Frame ordering is **not** validated here — that's
/// `framing::decode_message`'s job — so this helper tolerates
/// non-data-object frames in any phase and simply skips them.
///
/// Returns:
///   * `Ok(())` — every checked data-object frame's flag was set
///     and its slot matched the recomputed digest.
///   * `Err(MissingHash { object_index })` — at least one checked
///     frame had `HASH_PRESENT = 0`.
///   * `Err(HashMismatch { object_index: Some(i), .. })` — at
///     least one checked frame's slot disagreed with the
///     recomputed digest.
///   * `Err(Framing(_))` — buffer too short, preamble malformed,
///     `total_length` exceeds buffer, frame header malformed, or
///     frame extends past the message bound.
fn verify_data_object_frames(buf: &[u8], target_index: Option<usize>) -> Result<()> {
    use crate::wire::{
        FRAME_HEADER_SIZE, FRAME_MAGIC, FrameHeader, POSTAMBLE_SIZE, PREAMBLE_SIZE, Preamble,
    };

    if buf.len() < PREAMBLE_SIZE {
        return Err(TensogramError::Framing(format!(
            "buffer too short for preamble while verifying hashes: {} < {PREAMBLE_SIZE}",
            buf.len()
        )));
    }

    // Match `framing::decode_message`'s message-bounded walk:
    // when the preamble carries a non-zero `total_length` (the
    // back-filled / closed-stream case), restrict the verify
    // pre-pass to that range; otherwise fall back to "everything
    // up to the trailing postamble".  This is what makes
    // concatenated `[msgA][msgB]` buffers verify only msgA — the
    // same scope `decode` returns.
    let preamble = Preamble::read_from(buf)?;
    let msg_end = if preamble.total_length > 0 {
        let total_len = usize::try_from(preamble.total_length).map_err(|_| {
            TensogramError::Framing(format!(
                "preamble total_length ({}) overflows usize while verifying hashes",
                preamble.total_length
            ))
        })?;
        if total_len > buf.len() {
            return Err(TensogramError::Framing(format!(
                "preamble total_length ({total_len}) exceeds buffer size ({}) \
                 while verifying hashes",
                buf.len()
            )));
        }
        total_len.saturating_sub(POSTAMBLE_SIZE)
    } else {
        buf.len().checked_sub(POSTAMBLE_SIZE).ok_or_else(|| {
            TensogramError::Framing(format!(
                "buffer too short for postamble while verifying hashes: {} < {POSTAMBLE_SIZE}",
                buf.len()
            ))
        })?
    };

    let mut pos = PREAMBLE_SIZE;
    let mut object_index: usize = 0;
    while pos + FRAME_HEADER_SIZE <= msg_end {
        // Tolerate alignment padding / between-frame slop the same
        // way `framing::decode_message` does — advance one byte at
        // a time until we land on a frame magic.  Pathological
        // inputs that contain `b"FR"` mid-padding would be caught
        // by the subsequent `FrameHeader::read_from` validation.
        if &buf[pos..pos + FRAME_MAGIC.len()] != FRAME_MAGIC {
            pos += 1;
            continue;
        }
        let fh = FrameHeader::read_from(&buf[pos..])?;
        // `try_from` rather than `as usize` so a 32-bit target
        // with a >4 GB frame fails cleanly instead of truncating
        // + producing a wrong hash slice.
        let frame_total = usize::try_from(fh.total_length).map_err(|_| {
            TensogramError::Framing(format!(
                "frame total_length ({}) overflows usize on this target \
                 while verifying hash at offset {pos}",
                fh.total_length
            ))
        })?;
        let frame_end = pos.checked_add(frame_total).ok_or_else(|| {
            TensogramError::Framing(format!(
                "frame total_length overflow at offset {pos} while verifying hashes"
            ))
        })?;
        // Bound by `msg_end` rather than `buf.len()`: a frame that
        // straddles the end of the first message into a following
        // concatenated message would mean the buffer is malformed,
        // since `framing::decode_message` has the same bound.
        if frame_end > msg_end {
            return Err(TensogramError::Framing(format!(
                "frame at offset {pos} runs past first-message end \
                 ({frame_end} > {msg_end}) while verifying hashes"
            )));
        }

        if fh.frame_type.is_data_object() {
            // Run the hash check unless a `target_index` was
            // supplied and this is not the targeted object —
            // either way, we still increment `object_index`
            // afterwards so subsequent frames see the right
            // counter.
            let should_check = target_index.is_none_or(|t| t == object_index);
            if should_check {
                let frame_bytes = &buf[pos..frame_end];
                match crate::hash::check_frame_hash(frame_bytes, fh.frame_type) {
                    Ok(true) => {}
                    Ok(false) => {
                        return Err(TensogramError::MissingHash { object_index });
                    }
                    Err(e) => {
                        return Err(crate::error::with_object_index(e, object_index));
                    }
                }
                // When verifying a single targeted object we can
                // stop walking once it's checked — no need to hash
                // the rest of the message.
                if target_index.is_some() {
                    return Ok(());
                }
            }
            object_index += 1;
        }

        // 8-byte alignment padding may follow `ENDF`; advance to
        // the next aligned boundary.
        pos = (frame_end + 7) & !7;
    }
    Ok(())
}

/// Decode all objects from a message buffer.
/// Returns (global_metadata, list of (descriptor, decoded_data)).
///
/// When `options.threads > 0` (or `TENSOGRAM_THREADS` is set),
/// per-object decode work is parallelised using the axis-B-first
/// policy documented in
/// `docs/src/guide/multi-threaded-pipeline.md`.  Output bytes are
/// byte-identical to the sequential path regardless of thread count.
#[tracing::instrument(skip(buf, options), fields(buf_len = buf.len()))]
pub fn decode(buf: &[u8], options: &DecodeOptions) -> Result<(GlobalMetadata, Vec<DecodedObject>)> {
    // Hash verification runs as a *pre-pass*: we walk every
    // data-object frame's inline xxh3 BEFORE
    // `framing::decode_message` parses any CBOR descriptor.
    // That means a tamper inside the CBOR-encoded region (which
    // IS covered by the frame body hash, per `WIRE_FORMAT.md`
    // §2.4) surfaces as `HashMismatch` rather than the
    // `Metadata` error a downstream CBOR parse would otherwise
    // produce.  The xxh3 work here is identical to a per-object
    // re-check after decoding — bytes are walked exactly once,
    // and the cost is dominated by the single body hash, so
    // verifying first costs effectively nothing extra.
    if options.verify_hash {
        verify_data_object_frames(buf, None)?;
    }

    let msg = framing::decode_message(buf)?;

    let budget = crate::parallel::resolve_budget(options.threads)?;
    let total_bytes: usize = msg.objects.iter().map(|(_, p, _, _)| p.len()).sum();
    let parallel =
        crate::parallel::should_parallelise(budget, total_bytes, options.parallel_threshold_bytes);
    let any_axis_b = msg.objects.iter().any(|(d, _, _, _)| {
        crate::parallel::is_axis_b_friendly(&d.encoding, &d.filter, &d.compression)
    });
    let use_axis_a = parallel && crate::parallel::use_axis_a(msg.objects.len(), budget, any_axis_b);
    let intra_codec_threads = if parallel && !use_axis_a { budget } else { 0 };

    let decode_one = |(desc, payload_bytes, mask_region, _offset): &(
        DataObjectDescriptor,
        &[u8],
        &[u8],
        usize,
    )|
     -> Result<DecodedObject> {
        let mut decoded = decode_single_object_with_backend(
            desc,
            payload_bytes,
            options,
            options.compression_backend,
            intra_codec_threads,
        )?;
        if options.restore_non_finite {
            crate::restore::restore_non_finite_into(
                &mut decoded,
                desc,
                mask_region,
                output_byte_order(desc, options),
            )?;
        }
        Ok((desc.clone(), decoded))
    };

    let data_objects: Vec<DecodedObject> = if use_axis_a {
        #[cfg(feature = "threads")]
        {
            use rayon::prelude::*;
            crate::parallel::with_pool(budget, || {
                msg.objects
                    .par_iter()
                    .map(&decode_one)
                    .collect::<Result<Vec<_>>>()
            })?
        }
        #[cfg(not(feature = "threads"))]
        {
            msg.objects.iter().map(decode_one).collect::<Result<_>>()?
        }
    } else {
        crate::parallel::run_maybe_pooled(budget, parallel, intra_codec_threads, || {
            msg.objects.iter().map(decode_one).collect::<Result<_>>()
        })?
    };

    Ok((msg.global_metadata, data_objects))
}

/// Decode only global metadata from a message buffer, skipping payloads.
pub fn decode_metadata(buf: &[u8]) -> Result<GlobalMetadata> {
    framing::decode_metadata_only(buf)
}

/// Decode all objects from a message buffer AND return the raw
/// decompressed bitmasks alongside the substituted payloads.
///
/// Like [`decode`], but the returned payloads always contain `0.0`
/// at non-finite positions — restoration is **not** applied.
/// Callers get the raw [`restore::DecodedMaskSet`] for each object
/// and can apply the masks manually (e.g. to convert to a
/// domain-specific missing-value representation, or to aggregate
/// missing-count statistics without materialising the canonical
/// NaN / Inf bytes).
///
/// See `docs/src/guide/nan-inf-handling.md` for the companion user
/// guide and `plans/WIRE_FORMAT.md` §6.5 for the wire-format details.
pub fn decode_with_masks(
    buf: &[u8],
    options: &DecodeOptions,
) -> Result<(GlobalMetadata, Vec<DecodedObjectWithMasks>)> {
    let msg = framing::decode_message(buf)?;

    let budget = crate::parallel::resolve_budget(options.threads)?;
    let total_bytes: usize = msg.objects.iter().map(|(_, p, _, _)| p.len()).sum();
    let parallel =
        crate::parallel::should_parallelise(budget, total_bytes, options.parallel_threshold_bytes);
    let intra_codec_threads = if parallel { budget } else { 0 };

    // Local options snapshot with restore_non_finite forced off —
    // this API returns masks alongside a 0-substituted payload by
    // design.
    let mut decode_opts = options.clone();
    decode_opts.restore_non_finite = false;

    let decode_one = |(desc, payload_bytes, mask_region, _offset): &(
        DataObjectDescriptor,
        &[u8],
        &[u8],
        usize,
    )|
     -> Result<DecodedObjectWithMasks> {
        let payload = decode_single_object_with_backend(
            desc,
            payload_bytes,
            &decode_opts,
            options.compression_backend,
            intra_codec_threads,
        )?;
        let masks = crate::restore::decode_mask_set(desc, mask_region)?;
        Ok(DecodedObjectWithMasks {
            descriptor: desc.clone(),
            payload,
            masks,
        })
    };

    // Advanced API — axis-A parallelism is not worth the complexity
    // here since this path is intended for niche workflows (custom
    // missing-value representations, missing-count stats).  Mainline
    // consumers use `decode()` with `restore_non_finite=true` which
    // does have the full axis-A/B dispatch.
    let objects: Vec<DecodedObjectWithMasks> =
        crate::parallel::run_maybe_pooled(budget, parallel, intra_codec_threads, || {
            msg.objects.iter().map(decode_one).collect::<Result<_>>()
        })?;

    Ok((msg.global_metadata, objects))
}

pub use crate::restore::{DecodedMaskSet, DecodedObjectWithMasks};

/// Decode global metadata **and** per-object descriptors without decoding
/// any payload data.
///
/// This is cheaper than [`decode`] because the pipeline (decompression,
/// filter reversal, endian swap) is never executed.  Use it when you only
/// need shapes, dtypes, and metadata — e.g. for building xarray Datasets
/// at open time.
pub fn decode_descriptors(buf: &[u8]) -> Result<(GlobalMetadata, Vec<DataObjectDescriptor>)> {
    let msg = framing::decode_message(buf)?;
    let descriptors = msg
        .objects
        .into_iter()
        .map(|(desc, _, _, _)| desc)
        .collect();
    Ok((msg.global_metadata, descriptors))
}

/// Decode a single object by index (O(1) access via index frame).
/// Returns (global_metadata, descriptor, decoded_data).
pub fn decode_object(
    buf: &[u8],
    index: usize,
    options: &DecodeOptions,
) -> Result<(GlobalMetadata, DataObjectDescriptor, Vec<u8>)> {
    // Hash verification runs as a *pre-pass*: walk frame headers
    // and hash-verify object `index` BEFORE
    // `framing::decode_message` parses any CBOR descriptor.  See
    // `verify_data_object_frames` for the rationale.  Passing
    // `Some(index)` short-circuits once that single object is
    // checked, so the cost is one body hash recompute + the
    // header walk up to it (microseconds for typical messages).
    if options.verify_hash {
        verify_data_object_frames(buf, Some(index))?;
    }

    let msg = framing::decode_message(buf)?;

    if index >= msg.objects.len() {
        return Err(TensogramError::Object(format!(
            "object index {} out of range (num_objects={})",
            index,
            msg.objects.len()
        )));
    }

    let (desc, payload_bytes, mask_region, _frame_offset) = &msg.objects[index];

    // Single-object decode: axis A is impossible — spend the entire
    // budget (if any) on the codec internally (axis B).
    let budget = crate::parallel::resolve_budget(options.threads)?;
    let parallel = crate::parallel::should_parallelise(
        budget,
        payload_bytes.len(),
        options.parallel_threshold_bytes,
    );
    let intra_codec_threads = if parallel { budget } else { 0 };

    let mut decoded =
        crate::parallel::run_maybe_pooled(budget, parallel, intra_codec_threads, || {
            decode_single_object_with_backend(
                desc,
                payload_bytes,
                options,
                options.compression_backend,
                intra_codec_threads,
            )
        })?;

    if options.restore_non_finite {
        crate::restore::restore_non_finite_into(
            &mut decoded,
            desc,
            mask_region,
            output_byte_order(desc, options),
        )?;
    }

    Ok((msg.global_metadata, desc.clone(), decoded))
}

/// Decode a single data-object frame, given its complete bytes.
///
/// Mirrors [`decode_object`] but takes the bytes of one frame rather
/// than a full message.  Used by callers that obtained a single frame
/// via range-fetching (e.g. the WASM layout helpers that let the
/// TypeScript wrapper issue per-object HTTP Range reads) or by the
/// remote backend once it has fetched one indexed frame.
///
/// Returns `(descriptor, decoded_bytes)`.  There is no
/// `GlobalMetadata` in the return value because a single frame does
/// not carry it — the caller is expected to have it from a separate
/// metadata or header-chunk fetch.
pub fn decode_object_from_frame(
    frame_bytes: &[u8],
    options: &DecodeOptions,
) -> Result<(DataObjectDescriptor, Vec<u8>)> {
    // Verify-first ordering when `verify_hash` is set.  The hash
    // covers the post-encoding payload + CBOR descriptor, so a
    // corruption inside the CBOR bytes would otherwise surface as
    // `MetadataError` from `decode_data_object_frame` *before* the
    // hash check fires — leaking implementation detail through the
    // error class.  Verifying ahead of the CBOR parse guarantees
    // every corruption inside the hashed region surfaces as
    // `HashMismatch`, matching the contract documented on
    // [`DecodeOptions::verify_hash`].
    //
    // At this layer we don't know the surrounding object index —
    // the caller has fetched a single frame's bytes by some prior
    // path — so the resulting `MissingHash` / `HashMismatch`
    // carries the placeholder `0`.  Callers that have a real index
    // (typically the remote indexed fast paths in `remote.rs`)
    // re-stamp it via [`crate::error::with_object_index`].
    if options.verify_hash {
        use crate::wire::{FRAME_HEADER_SIZE, FrameHeader};
        let fh = FrameHeader::read_from(frame_bytes)?;
        // Slice down to total_length in case the caller passed a
        // larger buffer (frame + trailing alignment); the hash
        // helper is strict about ENDF placement.  `try_from`
        // bails cleanly on a 32-bit target with a >4 GB frame
        // (exotic, but a silent `as usize` truncation here would
        // hash too few bytes and surface as a spurious mismatch).
        let total = usize::try_from(fh.total_length).map_err(|_| {
            TensogramError::Framing(format!(
                "frame total_length ({}) overflows usize on this target; \
                 cannot verify hash on decode_object_from_frame",
                fh.total_length
            ))
        })?;
        if total < FRAME_HEADER_SIZE || total > frame_bytes.len() {
            return Err(TensogramError::Framing(format!(
                "frame total_length ({}) outside bounds of supplied frame buffer \
                 ({} bytes); cannot verify hash on decode_object_from_frame",
                fh.total_length,
                frame_bytes.len()
            )));
        }
        if !crate::hash::check_frame_hash(&frame_bytes[..total], fh.frame_type)? {
            return Err(TensogramError::MissingHash { object_index: 0 });
        }
    }
    let (desc, payload_bytes, mask_region, _) = framing::decode_data_object_frame(frame_bytes)?;

    let budget = crate::parallel::resolve_budget(options.threads)?;
    let parallel = crate::parallel::should_parallelise(
        budget,
        payload_bytes.len(),
        options.parallel_threshold_bytes,
    );
    let intra_codec_threads = if parallel { budget } else { 0 };

    let mut decoded =
        crate::parallel::run_maybe_pooled(budget, parallel, intra_codec_threads, || {
            decode_single_object_with_backend(
                &desc,
                payload_bytes,
                options,
                options.compression_backend,
                intra_codec_threads,
            )
        })?;

    if options.restore_non_finite {
        crate::restore::restore_non_finite_into(
            &mut decoded,
            &desc,
            mask_region,
            output_byte_order(&desc, options),
        )?;
    }

    Ok((desc, decoded))
}

/// Decode partial ranges from a single data-object frame.
///
/// Companion to [`decode_object_from_frame`] for the
/// `decode_range` code path: takes one frame's bytes, extracts the
/// requested element ranges from its payload, and applies mask-aware
/// non-finite restoration per range.
///
/// Returns `(descriptor, parts)` in the same shape as
/// [`decode_range`].
///
/// **`options.verify_hash` is ignored** by this function.  The
/// inline hash covers the whole post-encoding payload of the
/// source frame; verifying it requires reading every byte that
/// range decode is designed to avoid.  When integrity matters,
/// call [`decode_object_from_frame`] with `verify_hash=true` —
/// that path materialises the full body anyway, so the
/// verification is free.  See `plans/DESIGN.md`
/// §"Integrity Hashing" and `plans/WIRE_FORMAT.md` §11.1.
pub fn decode_range_from_frame(
    frame_bytes: &[u8],
    ranges: &[(u64, u64)],
    options: &DecodeOptions,
) -> Result<(DataObjectDescriptor, Vec<Vec<u8>>)> {
    let (desc, payload_bytes, mask_region, _) = framing::decode_data_object_frame(frame_bytes)?;
    let mut parts = decode_range_from_payload(&desc, payload_bytes, ranges, options)?;
    if options.restore_non_finite && desc.masks.is_some() {
        let mask_set = crate::restore::decode_mask_set(&desc, mask_region)?;
        crate::restore::restore_non_finite_into_ranges(
            &mut parts,
            &desc,
            ranges,
            &mask_set,
            output_byte_order(&desc, options),
        )?;
    }
    Ok((desc, parts))
}

/// Decode partial ranges from a data object.
///
/// `ranges` is a list of (element_offset, element_count) pairs.
///
/// Returns `(descriptor, parts)` where `parts` contains one `Vec<u8>`
/// per range.  The descriptor is included so callers can determine
/// the dtype without a separate lookup.
///
/// **`options.verify_hash` is ignored** by this function — see
/// [`decode_range_from_frame`] for the rationale.  Callers that
/// need integrity should reach for [`decode_object`] with
/// `verify_hash=true`.
pub fn decode_range(
    buf: &[u8],
    object_index: usize,
    ranges: &[(u64, u64)],
    options: &DecodeOptions,
) -> Result<(DataObjectDescriptor, Vec<Vec<u8>>)> {
    let msg = framing::decode_message(buf)?;

    if object_index >= msg.objects.len() {
        return Err(TensogramError::Object(format!(
            "object index {} out of range (num_objects={})",
            object_index,
            msg.objects.len()
        )));
    }

    let (desc, payload_bytes, mask_region, _) = &msg.objects[object_index];
    let mut parts = decode_range_from_payload(desc, payload_bytes, ranges, options)?;
    // Apply mask-aware NaN / Inf restoration to each requested
    // sub-range (see `docs/src/guide/nan-inf-handling.md`).
    if options.restore_non_finite && desc.masks.is_some() {
        let mask_set = crate::restore::decode_mask_set(desc, mask_region)?;
        crate::restore::restore_non_finite_into_ranges(
            &mut parts,
            desc,
            ranges,
            &mask_set,
            output_byte_order(desc, options),
        )?;
    }
    Ok((desc.clone(), parts))
}

/// Byte order of the bytes coming out of `decode_pipeline` for
/// `desc`, given the caller's [`DecodeOptions`].  Used to tell
/// [`crate::restore`] which endianness to write canonical NaN / Inf
/// bit patterns in.
fn output_byte_order(
    desc: &DataObjectDescriptor,
    options: &DecodeOptions,
) -> tensogram_encodings::ByteOrder {
    if options.native_byte_order {
        tensogram_encodings::ByteOrder::native()
    } else {
        desc.byte_order
    }
}

pub fn decode_range_from_payload(
    desc: &DataObjectDescriptor,
    payload_bytes: &[u8],
    ranges: &[(u64, u64)],
    options: &DecodeOptions,
) -> Result<Vec<Vec<u8>>> {
    if desc.filter != "none" {
        return Err(TensogramError::Encoding(
            "decode_range is not supported when a filter (e.g. shuffle) is applied".to_string(),
        ));
    }

    if desc.dtype.byte_width() == 0 {
        return Err(TensogramError::Encoding(
            "partial range decode not supported for bitmask dtype".to_string(),
        ));
    }

    // v3: per-object hash lives in the frame footer's inline slot
    // (see `plans/WIRE_FORMAT.md` §2.4).  Partial-range decode runs
    // at the sub-object level and has no access to the containing
    // frame bytes, so hash verification at this layer is a no-op —
    // callers wanting frame-level integrity should use
    // `tensogram validate --checksum` or call
    // `hash::verify_frame_hash` on the full frame directly.

    let num_elements = desc.num_elements()?;
    // Thread-budget dispatch for range decode.
    //
    // Each range is an independent decode call; parallelism is natural
    // when the caller requests multiple ranges.  Axis B is always
    // preferred when there's only one range.
    let budget = crate::parallel::resolve_budget(options.threads)?;
    // Work is proportional to decoded output, not the input payload —
    // sum the requested counts × element byte width.
    let elem_bytes = desc.dtype.byte_width();
    let total_bytes: usize = ranges
        .iter()
        .map(|(_, c)| (*c as usize).saturating_mul(elem_bytes))
        .sum();
    let parallel =
        crate::parallel::should_parallelise(budget, total_bytes, options.parallel_threshold_bytes);
    let axis_b_friendly =
        crate::parallel::is_axis_b_friendly(&desc.encoding, &desc.filter, &desc.compression);
    let use_axis_a = parallel && crate::parallel::use_axis_a(ranges.len(), budget, axis_b_friendly);
    let intra_codec_threads = if parallel && !use_axis_a { budget } else { 0 };

    let config = build_pipeline_config_with_backend(
        desc,
        num_elements,
        desc.dtype,
        options.compression_backend,
        intra_codec_threads,
    )?;

    let block_offsets = if desc.compression == "szip" {
        extract_block_offsets(&desc.params)?
    } else {
        Vec::new()
    };

    let decode_one = |offset: u64, count: u64| -> Result<Vec<u8>> {
        pipeline::decode_range_pipeline(
            payload_bytes,
            &config,
            &block_offsets,
            offset,
            count,
            options.native_byte_order,
        )
        .map_err(|e| {
            TensogramError::Encoding(format!("range (offset={offset}, count={count}): {e}"))
        })
    };

    let run_seq = || -> Result<Vec<Vec<u8>>> {
        ranges
            .iter()
            .map(|&(offset, count)| decode_one(offset, count))
            .collect()
    };

    let results: Vec<Vec<u8>> = if use_axis_a {
        #[cfg(feature = "threads")]
        {
            use rayon::prelude::*;
            crate::parallel::with_pool(budget, || {
                ranges
                    .par_iter()
                    .map(|&(offset, count)| decode_one(offset, count))
                    .collect::<Result<Vec<_>>>()
            })?
        }
        #[cfg(not(feature = "threads"))]
        {
            run_seq()?
        }
    } else {
        crate::parallel::run_maybe_pooled(budget, parallel, intra_codec_threads, run_seq)?
    };

    Ok(results)
}

/// Decode a single object payload using the specified compression backend
/// and intra-codec thread budget.
///
/// `intra_codec_threads == 0` preserves the pre-threads behaviour.
fn decode_single_object_with_backend(
    desc: &DataObjectDescriptor,
    payload_bytes: &[u8],
    options: &DecodeOptions,
    backend: pipeline::CompressionBackend,
    intra_codec_threads: u32,
) -> Result<Vec<u8>> {
    // v3: hash verification lives at the frame layer (see the inline
    // slot in `plans/WIRE_FORMAT.md` §2.4).  Use `validate --checksum`
    // for a full integrity sweep or `hash::verify_frame_hash(frame_bytes, ft)`
    // for programmatic per-frame verification — the decode path itself
    // is a pure deserialisation.
    let num_elements = desc.num_elements()?;
    let config = build_pipeline_config_with_backend(
        desc,
        num_elements,
        desc.dtype,
        backend,
        intra_codec_threads,
    )?;
    let decoded = pipeline::decode_pipeline(payload_bytes, &config, options.native_byte_order)
        .map_err(|e| TensogramError::Encoding(e.to_string()))?;

    Ok(decoded)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::dtype::Dtype;
    use crate::encode::{EncodeOptions, encode};
    use crate::types::ByteOrder;
    use std::collections::BTreeMap;

    fn make_global_meta() -> GlobalMetadata {
        GlobalMetadata {
            extra: BTreeMap::new(),
            ..Default::default()
        }
    }

    fn make_descriptor(shape: Vec<u64>) -> DataObjectDescriptor {
        let strides = if shape.is_empty() {
            vec![]
        } else {
            let mut s = vec![1u64; shape.len()];
            for i in (0..shape.len() - 1).rev() {
                s[i] = s[i + 1] * shape[i + 1];
            }
            s
        };
        DataObjectDescriptor {
            obj_type: "ntensor".to_string(),
            ndim: shape.len() as u64,
            shape,
            strides,
            dtype: Dtype::Float32,
            byte_order: ByteOrder::native(),
            encoding: "none".to_string(),
            filter: "none".to_string(),
            compression: "none".to_string(),
            params: BTreeMap::new(),
            masks: None,
        }
    }

    // ── corrupt descriptor CBOR → decode error ───────────────────────────

    #[test]
    fn test_decode_corrupt_message_bytes() {
        // Completely invalid bytes — not a valid tensogram message
        let garbage = vec![0xDE, 0xAD, 0xBE, 0xEF, 0x00, 0x01, 0x02, 0x03];
        let result = decode(&garbage, &DecodeOptions::default());
        assert!(result.is_err(), "decoding garbage should fail");
    }

    #[test]
    fn test_decode_truncated_message() {
        // Encode a valid message then truncate it
        let meta = make_global_meta();
        let desc = make_descriptor(vec![4]);
        let data = vec![0u8; 16];
        let encoded = encode(&meta, &[(&desc, &data)], &EncodeOptions::default()).unwrap();

        // Truncate to half
        let truncated = &encoded[..encoded.len() / 2];
        let result = decode(truncated, &DecodeOptions::default());
        assert!(result.is_err(), "decoding truncated message should fail");
    }

    #[test]
    fn test_decode_corrupted_cbor_in_message() {
        // Encode a valid message then corrupt the metadata frame CBOR.
        // The metadata CBOR starts right after preamble (24 bytes) +
        // frame header (16 bytes). Aggressively corrupt that region.
        let meta = make_global_meta();
        let desc = make_descriptor(vec![4]);
        let data = vec![42u8; 16];
        let mut encoded = encode(&meta, &[(&desc, &data)], &EncodeOptions::default()).unwrap();

        // Preamble = 24 bytes, Frame header = 16 bytes => CBOR starts at 40
        let cbor_start = 40;
        let corrupt_end = (cbor_start + 30).min(encoded.len());
        for byte in &mut encoded[cbor_start..corrupt_end] {
            *byte = 0xFF;
        }

        let result = decode(&encoded, &DecodeOptions::default());
        // Should fail because CBOR metadata or frame structure is corrupted
        assert!(result.is_err(), "decoding corrupted CBOR should fail");
    }

    // ── object index out of range in decode_object ───────────────────────

    #[test]
    fn test_decode_object_index_out_of_range() {
        let meta = make_global_meta();
        let desc = make_descriptor(vec![4]);
        let data = vec![0u8; 16];
        let encoded = encode(&meta, &[(&desc, &data)], &EncodeOptions::default()).unwrap();

        // Only 1 object (index 0), request index 1
        let result = decode_object(&encoded, 1, &DecodeOptions::default());
        assert!(result.is_err());
        let msg = result.unwrap_err().to_string();
        assert!(
            msg.contains("out of range"),
            "expected 'out of range', got: {msg}"
        );

        // Request a very large index
        let result = decode_object(&encoded, 999, &DecodeOptions::default());
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("out of range"));
    }

    #[test]
    fn test_decode_object_valid_index() {
        let meta = make_global_meta();
        let desc0 = make_descriptor(vec![2]);
        let data0 = vec![10u8; 8];
        let desc1 = make_descriptor(vec![3]);
        let data1 = vec![20u8; 12];

        let encoded = encode(
            &meta,
            &[(&desc0, data0.as_slice()), (&desc1, data1.as_slice())],
            &EncodeOptions::default(),
        )
        .unwrap();

        // Access object 0
        let (_, ret_desc, ret_data) =
            decode_object(&encoded, 0, &DecodeOptions::default()).unwrap();
        assert_eq!(ret_desc.shape, vec![2]);
        assert_eq!(ret_data, data0);

        // Access object 1
        let (_, ret_desc, ret_data) =
            decode_object(&encoded, 1, &DecodeOptions::default()).unwrap();
        assert_eq!(ret_desc.shape, vec![3]);
        assert_eq!(ret_data, data1);
    }

    // ── decode_range invalid byte ranges ─────────────────────────────────

    #[test]
    fn test_decode_range_object_index_out_of_range() {
        let meta = make_global_meta();
        let desc = make_descriptor(vec![4]);
        let data = vec![0u8; 16];
        let encoded = encode(&meta, &[(&desc, &data)], &EncodeOptions::default()).unwrap();

        let result = decode_range(&encoded, 5, &[(0, 2)], &DecodeOptions::default());
        assert!(result.is_err());
        let msg = result.unwrap_err().to_string();
        assert!(
            msg.contains("out of range"),
            "expected 'out of range', got: {msg}"
        );
    }

    #[test]
    fn test_decode_range_exceeds_payload() {
        let meta = make_global_meta();
        let desc = make_descriptor(vec![4]); // 4 float32s = 16 bytes
        let data = vec![0u8; 16];
        let encoded = encode(&meta, &[(&desc, &data)], &EncodeOptions::default()).unwrap();

        // Request range offset=2, count=10 but only 4 elements
        let result = decode_range(&encoded, 0, &[(2, 10)], &DecodeOptions::default());
        assert!(result.is_err(), "range exceeding payload should fail");
    }

    #[test]
    fn test_decode_range_valid() {
        let meta = make_global_meta();
        let desc = make_descriptor(vec![8]); // 8 float32s = 32 bytes
        let data: Vec<u8> = (0..32).collect();
        let encoded = encode(&meta, &[(&desc, &data)], &EncodeOptions::default()).unwrap();

        let (ret_desc, parts) =
            decode_range(&encoded, 0, &[(0, 4)], &DecodeOptions::default()).unwrap();
        assert_eq!(ret_desc.shape, vec![8]);
        assert_eq!(parts.len(), 1);
        assert_eq!(parts[0].len(), 16); // 4 float32s = 16 bytes
    }

    #[test]
    fn test_decode_range_empty_ranges() {
        let meta = make_global_meta();
        let desc = make_descriptor(vec![4]);
        let data = vec![0u8; 16];
        let encoded = encode(&meta, &[(&desc, &data)], &EncodeOptions::default()).unwrap();

        let (_, parts) = decode_range(&encoded, 0, &[], &DecodeOptions::default()).unwrap();
        assert!(parts.is_empty());
    }

    // ── decode_metadata ──────────────────────────────────────────────────

    #[test]
    fn test_decode_metadata_valid() {
        let meta = make_global_meta();
        let desc = make_descriptor(vec![4]);
        let data = vec![0u8; 16];
        let encoded = encode(&meta, &[(&desc, &data)], &EncodeOptions::default()).unwrap();

        let _decoded_meta = decode_metadata(&encoded).unwrap();
    }

    #[test]
    fn test_decode_metadata_corrupt() {
        let garbage = vec![0xFF; 50];
        let result = decode_metadata(&garbage);
        assert!(result.is_err(), "decode_metadata on garbage should fail");
    }

    // ── decode_descriptors ───────────────────────────────────────────────

    #[test]
    fn test_decode_descriptors_valid() {
        let meta = make_global_meta();
        let desc0 = make_descriptor(vec![4]);
        let desc1 = make_descriptor(vec![2, 3]);
        let data0 = vec![0u8; 16];
        let data1 = vec![0u8; 24];
        let encoded = encode(
            &meta,
            &[(&desc0, data0.as_slice()), (&desc1, data1.as_slice())],
            &EncodeOptions::default(),
        )
        .unwrap();

        let (_decoded_meta, descs) = decode_descriptors(&encoded).unwrap();
        assert_eq!(descs.len(), 2);
        assert_eq!(descs[0].shape, vec![4]);
        assert_eq!(descs[1].shape, vec![2, 3]);
    }

    // ── decode_range with filter=shuffle → error ─────────────────────────

    #[test]
    fn test_decode_range_filter_shuffle_rejected() {
        let meta = make_global_meta();
        let mut desc = make_descriptor(vec![100]);
        desc.filter = "shuffle".to_string();
        desc.params.insert(
            "shuffle_element_size".to_string(),
            ciborium::Value::Integer(4.into()),
        );
        // Finite f32 data (avoid tripping the 0.17 default-reject
        // finite-check with byte-pattern NaN bits).
        let data: Vec<u8> = (0..100).flat_map(|i| (i as f32).to_ne_bytes()).collect();

        let encoded = encode(&meta, &[(&desc, &data)], &EncodeOptions::default()).unwrap();

        let result = decode_range(&encoded, 0, &[(0, 10)], &DecodeOptions::default());
        assert!(result.is_err());
        let msg = result.unwrap_err().to_string();
        assert!(
            msg.contains("filter") || msg.contains("shuffle"),
            "expected filter/shuffle error, got: {msg}"
        );
    }

    // ── decode_range with bitmask dtype → error ──────────────────────────

    #[test]
    fn test_decode_range_bitmask_dtype_rejected() {
        let meta = make_global_meta();
        let desc = DataObjectDescriptor {
            obj_type: "ntensor".to_string(),
            ndim: 1,
            shape: vec![16],
            strides: vec![1],
            dtype: Dtype::Bitmask,
            byte_order: ByteOrder::native(),
            encoding: "none".to_string(),
            filter: "none".to_string(),
            compression: "none".to_string(),
            params: BTreeMap::new(),
            masks: None,
        };
        let data = vec![0xFF; 2]; // ceil(16/8) = 2 bytes

        let encoded = encode(&meta, &[(&desc, &data)], &EncodeOptions::default()).unwrap();

        let result = decode_range(&encoded, 0, &[(0, 8)], &DecodeOptions::default());
        assert!(result.is_err());
        let msg = result.unwrap_err().to_string();
        assert!(
            msg.contains("bitmask"),
            "expected bitmask error, got: {msg}"
        );
    }

    // ── DecodeOptions defaults ───────────────────────────────────────────

    #[test]
    fn test_decode_options_defaults() {
        let opts = DecodeOptions::default();
        assert!(opts.native_byte_order);
        assert!(opts.restore_non_finite);
    }

    // ── decode with unknown encoding in descriptor ───────────────────────

    #[test]
    fn test_decode_unknown_encoding_in_descriptor() {
        // We need to craft a message with an unknown encoding.
        // Easiest: encode a valid message, then manually patch the CBOR
        // descriptor's encoding field. Instead, use build_pipeline_config directly.
        let mut desc = make_descriptor(vec![4]);
        desc.encoding = "foobar".to_string();

        let result = crate::encode::build_pipeline_config_with_backend(
            &desc,
            4,
            Dtype::Float32,
            pipeline::CompressionBackend::default(),
            0,
        );
        assert!(result.is_err());
        let msg = result.unwrap_err().to_string();
        assert!(
            msg.contains("unknown encoding"),
            "expected 'unknown encoding', got: {msg}"
        );
    }

    // ── decode with unknown compression in descriptor ────────────────────

    #[test]
    fn test_decode_unknown_compression_in_descriptor() {
        let mut desc = make_descriptor(vec![4]);
        desc.compression = "quantum_compress".to_string();

        let result = crate::encode::build_pipeline_config_with_backend(
            &desc,
            4,
            Dtype::Float32,
            pipeline::CompressionBackend::default(),
            0,
        );
        assert!(result.is_err());
        let msg = result.unwrap_err().to_string();
        assert!(
            msg.contains("unknown compression"),
            "expected 'unknown compression', got: {msg}"
        );
    }

    // ── extract_block_offsets error paths ─────────────────────────────────

    #[test]
    fn test_extract_block_offsets_missing() {
        let params = BTreeMap::new();
        let result = extract_block_offsets(&params);
        assert!(result.is_err());
        let msg = result.unwrap_err().to_string();
        assert!(
            msg.contains("szip_block_offsets"),
            "expected szip_block_offsets error, got: {msg}"
        );
    }

    #[test]
    fn test_extract_block_offsets_wrong_type() {
        let mut params = BTreeMap::new();
        params.insert(
            "szip_block_offsets".to_string(),
            ciborium::Value::Text("not an array".to_string()),
        );
        let result = extract_block_offsets(&params);
        assert!(result.is_err());
        let msg = result.unwrap_err().to_string();
        assert!(
            msg.contains("must be an array"),
            "expected 'must be an array', got: {msg}"
        );
    }

    #[test]
    fn test_extract_block_offsets_non_integer_elements() {
        let mut params = BTreeMap::new();
        params.insert(
            "szip_block_offsets".to_string(),
            ciborium::Value::Array(vec![
                ciborium::Value::Float(1.5), // not an integer
            ]),
        );
        let result = extract_block_offsets(&params);
        assert!(result.is_err());
        let msg = result.unwrap_err().to_string();
        assert!(
            msg.contains("integers"),
            "expected integers error, got: {msg}"
        );
    }

    #[test]
    fn test_extract_block_offsets_valid() {
        let mut params = BTreeMap::new();
        params.insert(
            "szip_block_offsets".to_string(),
            ciborium::Value::Array(vec![
                ciborium::Value::Integer(0.into()),
                ciborium::Value::Integer(100.into()),
                ciborium::Value::Integer(200.into()),
            ]),
        );
        let result = extract_block_offsets(&params).unwrap();
        assert_eq!(result, vec![0, 100, 200]);
    }
}
