//! Aligned static `OffsetMap` — Mode 2 of the two layout modes (ADR-002).
//!
//! Fields have fixed positions with natural alignment padding.
//! Variable-length fields get a 4-byte length prefix at a known offset;
//! the variable data is not included in the static layout. Used for
//! mmap-friendly formats (metatensor, safetensors).
//!
//! The offset computation is a recursive walk of the BAST typed tree
//! ([`crate::bast::BastDoc`]). Nested structs propagate field path
//! prefixes (producing dotted paths like `"header.version"`). Alignment
//! padding is inserted before each field to satisfy the field's alignment
//! requirement (natural alignment by default, overridable via the
//! `"align"` annotation).

use crate::bast::{
    BastArray, BastDefKind, BastDoc, BastField, BastStruct, BastType,
};
use crate::error::AlkTypeError;
use crate::schema::{AlkTypeKind, VariableEncoding};

/// A byte range within a buffer.
///
/// Produced by [`OffsetMap::compute`] for each field in a schema. The
/// range is half-open: `start..end`. `end - start` is the field's byte
/// size in the static layout (for variable-length fields, this is the
/// size of the length prefix, the `{offset, length}` pair, or the
/// `maxLength` reservation — not the variable data itself).
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ByteRange {
    /// Inclusive start byte offset.
    pub start: usize,
    /// Exclusive end byte offset.
    pub end: usize,
}

impl ByteRange {
    /// Byte length of the range (`end - start`).
    pub fn len(&self) -> usize {
        self.end - self.start
    }

    /// True if the range covers zero bytes.
    pub fn is_empty(&self) -> bool {
        self.end == self.start
    }
}

/// A flat table of `(field_path, byte_range)` pairs computed from a BAST
/// document.
///
/// Fields have fixed positions with natural alignment padding.
/// Used for mmap-friendly formats (metatensor, safetensors) where random
/// access by field path is required — the consumer can read field N
/// without reading fields `0..N-1` first.
///
/// Construct via [`OffsetMap::compute`]. Variable-length fields appear
/// in the table as their fixed-position portion only (length prefix,
/// `{offset, length}` pair, or `maxLength` reservation); the variable
/// data lives outside the static layout.
#[derive(Debug)]
pub struct OffsetMap {
    fields: Vec<(String, ByteRange)>,
    total_size: usize,
}

impl OffsetMap {
    /// Compute the offset map from a BAST document.
    ///
    /// Walks the BAST typed tree recursively, computing byte positions
    /// for each field based on type sizes, field order, and alignment.
    /// The root type must be a struct.
    ///
    /// # Errors
    ///
    /// Returns [`AlkTypeError::Schema`] if the root type is not a struct.
    /// Returns [`AlkTypeError::Offset`] for unsupported type combinations
    /// encountered during the walk (e.g. unions, which are rejected in
    /// aligned mode per ADR-008).
    pub fn compute<'a>(doc: &'a BastDoc<'a>) -> Result<Self, AlkTypeError> {
        let root_def = doc.root_def();
        let struct_node = match root_def.kind() {
            BastDefKind::Struct(s) => s,
            other => {
                return Err(AlkTypeError::Schema(format!(
                    "OffsetMap::compute requires a struct at the root, got {kind}",
                    kind = other.alk_kind()
                )));
            }
        };
        let mut ctx = ComputeCtx {
            doc,
            fields: Vec::new(),
            offset: 0,
        };
        let struct_default_align = struct_node.align().unwrap_or(1).max(1);
        let (total, _align) = ctx.compute_struct(struct_node, "", struct_default_align)?;
        Ok(Self {
            fields: ctx.fields,
            total_size: total,
        })
    }

    /// Look up a field's byte range by dotted path (e.g., `"header.version"`).
    ///
    /// Returns `None` if no field with the given path was recorded. For
    /// TUnion byte-offset discriminators, the discriminator is recorded
    /// under the synthetic path `"__discriminator"` (qualified by the
    /// union field's path, e.g. `"payload.__discriminator"`).
    pub fn get(&self, field_path: &str) -> Option<&ByteRange> {
        self.fields
            .iter()
            .find(|(path, _)| path == field_path)
            .map(|(_, range)| range)
    }

    /// The total size of the struct in bytes (including trailing alignment padding).
    pub fn total_size(&self) -> usize {
        self.total_size
    }

    /// Iterate over all `(field_path, byte_range)` pairs in insertion order.
    ///
    /// Field order matches the BAST `fields` array order. Nested struct
    /// fields appear after their parent's path prefix.
    pub fn iter(&self) -> impl Iterator<Item = &(String, ByteRange)> {
        self.fields.iter()
    }
}

/// Mutable context threaded through the recursive offset computation.
///
/// Carries the running `offset`, the accumulating `fields` vec, and a
/// reference to the BAST doc for `$ref` resolution. Grouping these
/// keeps the recursive helper signatures small.
struct ComputeCtx<'d> {
    doc: &'d BastDoc<'d>,
    fields: Vec<(String, ByteRange)>,
    offset: usize,
}

/// Result of laying out a single field: its alignment.
struct FieldLayout {
    align: usize,
}

impl<'d> ComputeCtx<'d> {
    /// Recurse into a `BastStruct`, appending `(field_path, ByteRange)`
    /// pairs to `self.fields` and advancing `self.offset`.
    ///
    /// Returns `(total_size, alignment)` where `total_size` includes
    /// trailing alignment padding and `alignment` is the struct's
    /// effective alignment (its own `align` annotation, or the max of its
    /// fields' alignments).
    ///
    /// `prefix` is the dotted path prefix for nested fields (empty at the
    /// top level). `parent_struct_align` is the default alignment a field
    /// inherits when it specifies neither its own `align` annotation nor
    /// a natural alignment larger than the default.
    fn compute_struct(
        &mut self,
        struct_node: &BastStruct<'d>,
        prefix: &str,
        parent_struct_align: usize,
    ) -> Result<(usize, usize), AlkTypeError> {
        let struct_default_align = struct_node.align().unwrap_or(parent_struct_align).max(1);
        let mut max_align: usize = 1;
        let struct_start = self.offset;
        let fields = struct_node.fields();
        let field_count = fields.len();
        for (i, field) in fields.iter().enumerate() {
            let field_path = if prefix.is_empty() {
                field.name().to_string()
            } else {
                format!("{prefix}.{}", field.name())
            };
            if i < field_count - 1 {
                if let Some(kind) = field_variable_kind(field) {
                    let encoding = field.encoding();
                    let max_length = field.max_length();
                    let is_inline_length_prefixed =
                        encoding == VariableEncoding::LengthPrefixed && max_length.is_none();
                    if is_inline_length_prefixed {
                        return Err(AlkTypeError::Offset {
                            field_path: field_path.clone(),
                            reason: format!(
                                "non-final inline length-prefixed variable field \
                                 ({kind}) in aligned mode: the variable data would \
                                 clobber subsequent fields. Use `maxLength` \
                                 (fixed-size reservation) or \
                                 `\"encoding\": \"offset-indirect\"`, or move this \
                                 field to the last position in the struct. (ADR-006)"
                            ),
                        });
                    }
                }
            }
            let layout = self.compute_field(field, &field_path, struct_default_align)?;
            if layout.align > max_align {
                max_align = layout.align;
            }
        }

        let effective_align = struct_node.align().unwrap_or(max_align).max(1);
        align_up(&mut self.offset, effective_align);
        let total = self.offset - struct_start;
        Ok((total, effective_align))
    }

    /// Compute the layout for a single field, advancing `self.offset`
    /// and appending any field paths to `self.fields`.
    fn compute_field(
        &mut self,
        field: &BastField<'d>,
        field_path: &str,
        struct_default_align: usize,
    ) -> Result<FieldLayout, AlkTypeError> {
        let ty = field.ty();
        let resolved = self.doc.resolve_typeref(ty)?;
        match &resolved {
            BastType::Struct(s) => self.compute_struct_field(s, field, field_path, struct_default_align),
            BastType::Union(_) => Err(AlkTypeError::Offset {
                field_path: field_path.to_string(),
                reason: "TUnion is not supported in aligned static mode (ADR-008). \
                         Unions are the protocol dispatch pattern — use packed sequential \
                         mode (LayoutMode::Packed) for TUnion fields, or restructure as \
                         a struct with an explicit discriminator field."
                    .to_string(),
            }),
            BastType::Array(a) => self.compute_array_field(a, field, field_path, struct_default_align),
            BastType::Record(_) => {
                self.compute_variable_field(field, field_path, struct_default_align)
            }
            BastType::Primitive(k) if k.is_variable_length() => {
                self.compute_variable_field(field, field_path, struct_default_align)
            }
            BastType::Primitive(k) => {
                self.compute_fixed_field(*k, field, field_path, struct_default_align)
            }
            BastType::Enum(_) => self.compute_fixed_field(
                AlkTypeKind::Enum,
                field,
                field_path,
                struct_default_align,
            ),
            BastType::Ref(_) => Err(AlkTypeError::Offset {
                field_path: field_path.to_string(),
                reason: "internal: resolve_typeref returned a Ref".to_string(),
            }),
        }
    }

    /// Compute the layout for a fixed-size primitive field.
    fn compute_fixed_field(
        &mut self,
        kind: AlkTypeKind,
        field: &BastField<'d>,
        field_path: &str,
        struct_default_align: usize,
    ) -> Result<FieldLayout, AlkTypeError> {
        let size = kind.type_size().ok_or_else(|| AlkTypeError::Offset {
            field_path: field_path.to_string(),
            reason: format!("type_size returned None for fixed kind {kind}"),
        })?;
        let natural = kind.natural_alignment();
        let align = field_alignment(field, struct_default_align, natural);
        align_up(&mut self.offset, align);
        let start = self.offset;
        self.offset = start
            .checked_add(size)
            .ok_or_else(|| AlkTypeError::Offset {
                field_path: field_path.to_string(),
                reason: format!("offset {start} + size {size} overflows usize"),
            })?;
        self.push(field_path, start, start + size);
        Ok(FieldLayout { align })
    }

    /// Compute the layout for a nested struct field.
    ///
    /// Probes the nested struct's layout at a temporary offset of 0 to
    /// determine its total size and alignment, aligns the parent offset,
    /// then shifts the nested fields to their final positions.
    fn compute_struct_field(
        &mut self,
        struct_node: &BastStruct<'d>,
        field: &BastField<'d>,
        field_path: &str,
        struct_default_align: usize,
    ) -> Result<FieldLayout, AlkTypeError> {
        let inner_parent_align = field.align().unwrap_or(struct_default_align);
        let mut probe = ComputeCtx {
            doc: self.doc,
            fields: Vec::new(),
            offset: 0,
        };
        let (inner_total, inner_align) =
            probe.compute_struct(struct_node, field_path, inner_parent_align)?;

        let natural = inner_align;
        let align = field_alignment(field, struct_default_align, natural);
        align_up(&mut self.offset, align);
        let struct_start = self.offset;
        for (path, range) in probe.fields {
            self.fields.push((
                path,
                ByteRange {
                    start: struct_start + range.start,
                    end: struct_start + range.end,
                },
            ));
        }
        self.offset = struct_start
            .checked_add(inner_total)
            .ok_or_else(|| AlkTypeError::Offset {
                field_path: field_path.to_string(),
                reason: format!("struct start {struct_start} + inner total {inner_total} overflows usize"),
            })?;
        Ok(FieldLayout { align })
    }

    /// Compute the layout for a `BastArray` field.
    fn compute_array_field(
        &mut self,
        array: &BastArray<'d>,
        field: &BastField<'d>,
        field_path: &str,
        struct_default_align: usize,
    ) -> Result<FieldLayout, AlkTypeError> {
        let element_ty = array.element();
        let resolved_elem = self.doc.resolve_typeref(element_ty)?;
        let elem_kind = resolved_elem.alk_kind();
        if !elem_kind.is_fixed_size() {
            return Err(AlkTypeError::Offset {
                field_path: field_path.to_string(),
                reason: format!(
                    "array of variable-length element kind {elem_kind} is not supported (OQ-001)"
                ),
            });
        }

        let elem_size = elem_kind.type_size().ok_or_else(|| AlkTypeError::Offset {
            field_path: field_path.to_string(),
            reason: format!("element kind {elem_kind} has no fixed size"),
        })?;
        let elem_natural = elem_kind.natural_alignment();
        let elem_align = element_alignment(&resolved_elem, struct_default_align, elem_natural);
        let stride = round_up(elem_size, elem_align);
        let count = array.count();

        let array_align = field_alignment(field, struct_default_align, elem_align);

        align_up(&mut self.offset, array_align);
        let start = self.offset;
        for i in 0..count {
            let elem_start = start
                .checked_add(
                    i.checked_mul(stride)
                        .ok_or_else(|| AlkTypeError::Offset {
                            field_path: field_path.to_string(),
                            reason: format!("element index {i} × stride {stride} overflows usize"),
                        })?,
                )
                .ok_or_else(|| AlkTypeError::Offset {
                    field_path: field_path.to_string(),
                    reason: format!("element offset {start} + {i}×{stride} overflows usize"),
                })?;
            let elem_end = elem_start
                .checked_add(elem_size)
                .ok_or_else(|| AlkTypeError::Offset {
                    field_path: field_path.to_string(),
                    reason: format!("element end {elem_start} + {elem_size} overflows usize"),
                })?;
            let elem_path = format!("{field_path}[{i}]");
            self.push(&elem_path, elem_start, elem_end);
        }
        let array_size = count
            .checked_mul(stride)
            .ok_or_else(|| AlkTypeError::Offset {
                field_path: field_path.to_string(),
                reason: format!("array size {count} × stride {stride} overflows usize"),
            })?;
        self.offset = start
            .checked_add(array_size)
            .ok_or_else(|| AlkTypeError::Offset {
                field_path: field_path.to_string(),
                reason: format!("offset {start} + array size {array_size} overflows usize"),
            })?;
        Ok(FieldLayout { align: array_align })
    }

    /// Compute the layout for a variable-length field (String/Bytes/Record/Timestamp).
    ///
    /// In aligned static mode, three strategies are supported:
    /// - `maxLength` reservation: `maxLength` bytes at a fixed offset.
    /// - `offset-indirect` encoding: an 8-byte `{offset: u32, length: u32}` pair.
    /// - inline length-prefixing (default): a 4-byte length prefix.
    fn compute_variable_field(
        &mut self,
        field: &BastField<'d>,
        field_path: &str,
        struct_default_align: usize,
    ) -> Result<FieldLayout, AlkTypeError> {
        let encoding = field.encoding();
        let max_length = field.max_length();

        let (size, natural) = match (max_length, encoding) {
            (Some(max_len), _) => (max_len, 1),
            (None, VariableEncoding::OffsetIndirect) => (8, 4),
            (None, VariableEncoding::LengthPrefixed) => (4, 4),
        };

        let align = field_alignment(field, struct_default_align, natural);
        align_up(&mut self.offset, align);
        let start = self.offset;
        self.offset = start
            .checked_add(size)
            .ok_or_else(|| AlkTypeError::Offset {
                field_path: field_path.to_string(),
                reason: format!("offset {start} + size {size} overflows usize"),
            })?;
        self.push(field_path, start, start + size);
        Ok(FieldLayout { align })
    }

    /// Push a `(field_path, ByteRange)` pair onto the fields vec.
    fn push(&mut self, path: &str, start: usize, end: usize) {
        self.fields
            .push((path.to_string(), ByteRange { start, end }));
    }
}

/// If the field's type is a variable-length primitive, return its kind.
fn field_variable_kind(field: &BastField<'_>) -> Option<AlkTypeKind> {
    match field.ty() {
        BastType::Primitive(k) if k.is_variable_length() => Some(*k),
        _ => None,
    }
}

/// Resolve the field's alignment: field-level `align` annotation,
/// then the struct default, then the natural alignment.
fn field_alignment(field: &BastField<'_>, struct_default_align: usize, natural: usize) -> usize {
    if let Some(a) = field.align() {
        return a.max(1);
    }
    struct_default_align.max(natural).max(1)
}

/// Resolve an element type's alignment. Inline struct/array elements use
/// natural alignment 1; primitives use their natural alignment; field-level
/// `align` from the enclosing field node doesn't apply to inline element
/// schemas (BAST field-level annotations live on the field, not the
/// element TypeRef), so we fall back to the struct default.
fn element_alignment(
    elem_ty: &BastType<'_>,
    struct_default_align: usize,
    natural: usize,
) -> usize {
    match elem_ty {
        BastType::Struct(_) | BastType::Union(_) | BastType::Array(_) => {
            struct_default_align.max(1)
        }
        _ => struct_default_align.max(natural).max(1),
    }
}

/// Round `offset` up to the next multiple of `align`. No-op if `align <= 1`.
fn align_up(offset: &mut usize, align: usize) {
    if align <= 1 {
        return;
    }
    let rem = *offset % align;
    if rem != 0 {
        *offset += align - rem;
    }
}

/// Round `n` up to the next multiple of `align`.
fn round_up(n: usize, align: usize) -> usize {
    if align <= 1 {
        return n;
    }
    let rem = n % align;
    if rem == 0 {
        n
    } else {
        n + align - rem
    }
}

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

    fn map(root: &serde_json::Value, name: &str) -> OffsetMap {
        let doc = BastDoc::new(root, name).expect("bast doc");
        OffsetMap::compute(&doc).expect("offset map computation")
    }

    #[test]
    fn simple_fixed_fields_natural_alignment() {
        let root = json!({
            "$defs": {
                "S": {
                    "kind": "struct",
                    "fields": [
                        { "name": "flag", "kind": "uint8" },
                        { "name": "id",   "kind": "uint32" }
                    ]
                }
            }
        });
        let m = map(&root, "S");
        assert_eq!(m.get("flag"), Some(&ByteRange { start: 0, end: 1 }));
        assert_eq!(m.get("id"), Some(&ByteRange { start: 4, end: 8 }));
        assert_eq!(m.total_size(), 8);
    }

    #[test]
    fn u8_then_u32_three_bytes_padding() {
        let root = json!({
            "$defs": {
                "S": {
                    "kind": "struct",
                    "fields": [
                        { "name": "a", "kind": "uint8" },
                        { "name": "b", "kind": "uint32" }
                    ]
                }
            }
        });
        let m = map(&root, "S");
        assert_eq!(m.get("a"), Some(&ByteRange { start: 0, end: 1 }));
        assert_eq!(m.get("b"), Some(&ByteRange { start: 4, end: 8 }));
    }

    #[test]
    fn nested_struct_dotted_paths() {
        let root = json!({
            "$defs": {
                "S": {
                    "kind": "struct",
                    "fields": [
                        {
                            "name": "header",
                            "kind": {
                                "kind": "struct",
                                "fields": [
                                    { "name": "magic",   "kind": "uint32" },
                                    { "name": "version", "kind": "uint8" }
                                ]
                            }
                        },
                        { "name": "body", "kind": "uint32" }
                    ]
                }
            }
        });
        let m = map(&root, "S");
        assert_eq!(m.get("header.magic"), Some(&ByteRange { start: 0, end: 4 }));
        assert_eq!(
            m.get("header.version"),
            Some(&ByteRange { start: 4, end: 5 })
        );
        assert_eq!(m.get("body"), Some(&ByteRange { start: 8, end: 12 }));
        assert_eq!(m.total_size(), 12);
    }

    #[test]
    fn array_fixed_count_element_offsets() {
        let root = json!({
            "$defs": {
                "S": {
                    "kind": "struct",
                    "fields": [
                        { "name": "vals", "kind": { "kind": "array", "element": "uint32", "count": 3 } }
                    ]
                }
            }
        });
        let m = map(&root, "S");
        assert_eq!(m.get("vals[0]"), Some(&ByteRange { start: 0, end: 4 }));
        assert_eq!(m.get("vals[1]"), Some(&ByteRange { start: 4, end: 8 }));
        assert_eq!(m.get("vals[2]"), Some(&ByteRange { start: 8, end: 12 }));
        assert_eq!(m.total_size(), 12);
    }

    #[test]
    fn variable_string_length_prefix_at_known_offset() {
        let root = json!({
            "$defs": {
                "S": {
                    "kind": "struct",
                    "fields": [
                        { "name": "id",   "kind": "uint32" },
                        { "name": "name", "kind": "string" }
                    ]
                }
            }
        });
        let m = map(&root, "S");
        assert_eq!(m.get("id"), Some(&ByteRange { start: 0, end: 4 }));
        assert_eq!(m.get("name"), Some(&ByteRange { start: 4, end: 8 }));
        assert_eq!(m.total_size(), 8);
    }

    #[test]
    fn variable_string_max_length_reservation() {
        let root = json!({
            "$defs": {
                "S": {
                    "kind": "struct",
                    "fields": [
                        { "name": "id",   "kind": "uint32" },
                        { "name": "name", "kind": "string", "maxLength": 256 }
                    ]
                }
            }
        });
        let m = map(&root, "S");
        assert_eq!(m.get("id"), Some(&ByteRange { start: 0, end: 4 }));
        assert_eq!(m.get("name"), Some(&ByteRange { start: 4, end: 260 }));
        assert_eq!(m.total_size(), 260);
    }

    #[test]
    fn variable_string_offset_indirect_eight_bytes() {
        let root = json!({
            "$defs": {
                "S": {
                    "kind": "struct",
                    "fields": [
                        { "name": "id",   "kind": "uint32" },
                        { "name": "blob", "kind": "string", "encoding": "offset-indirect" }
                    ]
                }
            }
        });
        let m = map(&root, "S");
        assert_eq!(m.get("id"), Some(&ByteRange { start: 0, end: 4 }));
        assert_eq!(m.get("blob"), Some(&ByteRange { start: 4, end: 12 }));
        assert_eq!(m.total_size(), 12);
    }

    #[test]
    fn union_byte_discriminator_rejected_in_aligned_mode() {
        let root = json!({
            "$defs": {
                "S": {
                    "kind": "struct",
                    "fields": [
                        {
                            "name": "payload",
                            "kind": { "$ref": "#/$defs/Packet" }
                        }
                    ]
                },
                "Packet": {
                    "kind": "union",
                    "discriminator": { "kind": "byte", "offset": 0, "type": "uint8" },
                    "mapping": {
                        "5": { "$ref": "#/$defs/Read" },
                        "6": { "$ref": "#/$defs/Write" }
                    }
                },
                "Read": {
                    "kind": "struct",
                    "fields": [
                        { "name": "handle", "kind": "uint32" },
                        { "name": "length", "kind": "uint32" }
                    ]
                },
                "Write": {
                    "kind": "struct",
                    "fields": [
                        { "name": "handle", "kind": "uint32" },
                        { "name": "length", "kind": "uint32" },
                        { "name": "data",   "kind": "uint32" }
                    ]
                }
            }
        });
        let doc = BastDoc::new(&root, "S").expect("doc");
        let err = OffsetMap::compute(&doc).unwrap_err();
        assert!(matches!(err, AlkTypeError::Offset { .. }), "got {err:?}");
        let reason = match err {
            AlkTypeError::Offset { reason, .. } => reason,
            _ => unreachable!(),
        };
        assert!(reason.contains("ADR-008"), "reason: {reason}");
    }

    #[test]
    fn non_final_inline_string_rejected_in_aligned_mode() {
        let root = json!({
            "$defs": {
                "S": {
                    "kind": "struct",
                    "fields": [
                        { "name": "name", "kind": "string" },
                        { "name": "id",   "kind": "uint32" }
                    ]
                }
            }
        });
        let doc = BastDoc::new(&root, "S").expect("doc");
        let err = OffsetMap::compute(&doc).unwrap_err();
        match err {
            AlkTypeError::Offset { field_path, reason } => {
                assert_eq!(field_path, "name");
                assert!(reason.contains("ADR-006"), "reason: {reason}");
            }
            other => panic!("expected Offset, got {other:?}"),
        }
    }

    #[test]
    fn final_inline_string_allowed_in_aligned_mode() {
        let root = json!({
            "$defs": {
                "S": {
                    "kind": "struct",
                    "fields": [
                        { "name": "id",   "kind": "uint32" },
                        { "name": "name", "kind": "string" }
                    ]
                }
            }
        });
        let m = map(&root, "S");
        assert_eq!(m.get("id"), Some(&ByteRange { start: 0, end: 4 }));
        assert_eq!(m.get("name"), Some(&ByteRange { start: 4, end: 8 }));
    }

    #[test]
    fn non_final_maxlength_string_allowed_in_aligned_mode() {
        let root = json!({
            "$defs": {
                "S": {
                    "kind": "struct",
                    "fields": [
                        { "name": "name", "kind": "string", "maxLength": 256 },
                        { "name": "id",   "kind": "uint32" }
                    ]
                }
            }
        });
        let m = map(&root, "S");
        assert_eq!(m.get("name"), Some(&ByteRange { start: 0, end: 256 }));
        assert_eq!(m.get("id"), Some(&ByteRange { start: 256, end: 260 }));
    }

    #[test]
    fn non_final_offset_indirect_string_allowed_in_aligned_mode() {
        let root = json!({
            "$defs": {
                "S": {
                    "kind": "struct",
                    "fields": [
                        { "name": "blob", "kind": "string", "encoding": "offset-indirect" },
                        { "name": "id",   "kind": "uint32" }
                    ]
                }
            }
        });
        let m = map(&root, "S");
        assert_eq!(m.get("blob"), Some(&ByteRange { start: 0, end: 8 }));
        assert_eq!(m.get("id"), Some(&ByteRange { start: 8, end: 12 }));
    }

    #[test]
    fn struct_level_align_rounds_up_total() {
        let root = json!({
            "$defs": {
                "S": {
                    "kind": "struct",
                    "align": 16,
                    "fields": [
                        { "name": "flag", "kind": "uint8" }
                    ]
                }
            }
        });
        let m = map(&root, "S");
        assert_eq!(m.get("flag"), Some(&ByteRange { start: 0, end: 1 }));
        assert_eq!(m.total_size(), 16);
    }

    #[test]
    fn field_level_align_overrides_struct_default() {
        let root = json!({
            "$defs": {
                "S": {
                    "kind": "struct",
                    "align": 1,
                    "fields": [
                        { "name": "tag",  "kind": "uint8" },
                        { "name": "flag", "kind": "uint8", "align": 16 },
                        { "name": "id",   "kind": "uint32" }
                    ]
                }
            }
        });
        let m = map(&root, "S");
        assert_eq!(m.get("tag"), Some(&ByteRange { start: 0, end: 1 }));
        assert_eq!(m.get("flag"), Some(&ByteRange { start: 16, end: 17 }));
        assert_eq!(m.get("id"), Some(&ByteRange { start: 20, end: 24 }));
        assert_eq!(m.total_size(), 24);
    }

    #[test]
    fn field_align_smaller_than_struct_default() {
        let root = json!({
            "$defs": {
                "S": {
                    "kind": "struct",
                    "align": 8,
                    "fields": [
                        { "name": "a", "kind": "uint8" },
                        { "name": "b", "kind": "uint32", "align": 1 }
                    ]
                }
            }
        });
        let m = map(&root, "S");
        assert_eq!(m.get("a"), Some(&ByteRange { start: 0, end: 1 }));
        assert_eq!(m.get("b"), Some(&ByteRange { start: 1, end: 5 }));
        assert_eq!(m.total_size(), 8);
    }

    #[test]
    fn iter_returns_all_paths_in_order() {
        let root = json!({
            "$defs": {
                "S": {
                    "kind": "struct",
                    "fields": [
                        { "name": "a", "kind": "uint8" },
                        { "name": "b", "kind": "uint32" }
                    ]
                }
            }
        });
        let m = map(&root, "S");
        let paths: Vec<&String> = m.iter().map(|(p, _)| p).collect();
        assert_eq!(paths, vec!["a", "b"]);
    }

    #[test]
    fn compute_rejects_non_struct_top_level() {
        let root = json!({
            "$defs": {
                "U": {
                    "kind": "union",
                    "discriminator": { "kind": "byte", "offset": 0, "type": "uint8" },
                    "mapping": { "1": { "$ref": "#/$defs/A" } }
                },
                "A": { "kind": "struct", "fields": [] }
            }
        });
        let doc = BastDoc::new(&root, "U").expect("doc");
        let err = OffsetMap::compute(&doc).unwrap_err();
        assert!(matches!(err, AlkTypeError::Schema(_)));
    }

    #[test]
    fn byte_range_len_and_is_empty() {
        let r = ByteRange { start: 4, end: 8 };
        assert_eq!(r.len(), 4);
        assert!(!r.is_empty());
        let empty = ByteRange { start: 5, end: 5 };
        assert_eq!(empty.len(), 0);
        assert!(empty.is_empty());
    }
}