mig_assembly/

assembler.rs

//! Recursive descent assembler — MIG-guided segment consumption.
//!
//! The assembler walks the MIG tree structure and consumes matching
//! segments from the input. It produces a generic tree representation
//! that can be converted to typed PID structs.

use crate::cursor::SegmentCursor;
use crate::diagnostic::{StructureDiagnostic, StructureDiagnosticKind};
use crate::matcher;
use crate::tokenize::OwnedSegment;
use crate::AssemblyError;
use mig_types::schema::mig::{MigSchema, MigSegment, MigSegmentGroup};
use serde::{Deserialize, Serialize};

/// A generic assembled tree node (before PID-specific typing).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AssembledTree {
    pub segments: Vec<AssembledSegment>,
    pub groups: Vec<AssembledGroup>,
    /// Index in `segments` where post-group segments start (e.g., UNT, UNZ).
    /// Segments before this index appear before groups in EDIFACT order.
    #[serde(default)]
    pub post_group_start: usize,
    /// Root segments consumed between groups during assembly (e.g., UNS
    /// section separator in MSCONS). Key = index into `groups` vec; value =
    /// segments that appear immediately before that group in the EDIFACT
    /// stream. Empty for messages without inter-group root segments.
    #[serde(default, skip_serializing_if = "std::collections::BTreeMap::is_empty")]
    pub inter_group_segments: std::collections::BTreeMap<usize, Vec<AssembledSegment>>,
}

/// An assembled segment with its data elements.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AssembledSegment {
    pub tag: String,
    /// `elements[i][j]` = component `j` of element `i`
    pub elements: Vec<Vec<String>>,
    /// MIG `Number` attribute identifying this segment variant.
    /// Two segments with the same tag (e.g., DTM) but different roles
    /// (DTM+92 vs DTM+93) have distinct MIG numbers.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub mig_number: Option<String>,
}

/// An assembled segment group (may repeat).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AssembledGroup {
    pub group_id: String,
    pub repetitions: Vec<AssembledGroupInstance>,
}

/// One repetition of a segment group.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AssembledGroupInstance {
    pub segments: Vec<AssembledSegment>,
    pub child_groups: Vec<AssembledGroup>,
    /// MIG `Number` of the entry segment that identified this group instance's variant.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub entry_mig_number: Option<String>,
    /// All MIG `Number`s defined for this group variant — includes segments that
    /// may be absent in the EDIFACT but are defined in the MIG for this variant.
    ///
    /// Used by the validator to determine which AHB rules belong to this instance:
    /// a rule with `mig_number` in this set applies here, even if the segment is
    /// missing (which is then a missing-field error). Without this, rules for
    /// absent-but-required segments would be incorrectly filtered out.
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub variant_mig_numbers: Vec<String>,
    /// Segments that were present in the EDIFACT input but not defined in
    /// the PID-filtered MIG for this group. Only populated when the assembler
    /// runs with [`AssemblerConfig::skip_unknown_segments`] enabled.
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub skipped_segments: Vec<AssembledSegment>,
}

impl AssembledGroupInstance {
    /// Create a virtual `AssembledTree` scoped to this group instance.
    ///
    /// The instance's own segments become the tree's root segments,
    /// and its child groups become the tree's groups. This enables
    /// running `MappingEngine::map_all_forward()` on a single
    /// transaction group as if it were a complete message.
    pub fn as_assembled_tree(&self) -> AssembledTree {
        AssembledTree {
            segments: self.segments.clone(),
            groups: self.child_groups.clone(),
            post_group_start: self.segments.len(),
            inter_group_segments: std::collections::BTreeMap::new(),
        }
    }
}

/// Configuration for the assembler.
#[derive(Debug, Clone, Default)]
pub struct AssemblerConfig {
    /// When `true`, the assembler skips segments inside a group instance that
    /// don't match any remaining MIG slot, nested-group entry, or the group's
    /// entry tag (next repetition). Skipped segments are preserved on
    /// [`AssembledGroupInstance::skipped_segments`] for roundtrip re-emission.
    ///
    /// Default: `false` (strict AHB — unknown segments stall the cursor).
    pub skip_unknown_segments: bool,

    /// Qualifier-aware assembly: maps MIG `Number` to `(element_index, component_index, expected_value)`.
    ///
    /// When a bounded slot has a `number` with an entry in this map,
    /// `try_consume_segment` checks the input segment's value at the
    /// specified position. If it doesn't match, the slot is skipped (segment
    /// is for a different qualifier variant).
    ///
    /// Build from the PID schema JSON, or construct manually:
    /// `{ "00023" => (0, 0, "92".to_string()), "00024" => (0, 0, "93".to_string()) }`.
    ///
    /// Default: empty (positional assembly, no qualifier checking).
    pub qualifier_map: std::collections::HashMap<String, (usize, usize, String)>,

    /// Reject a slot when the input segment has a value outside the slot's
    /// allowed codes at any code-bearing position (not just the primary
    /// qualifier). Disambiguates PID-filtered slots that share a primary
    /// qualifier but differ on a secondary code (e.g. PID 55035 SG8/PIA
    /// variants all use 4347='5' but differ at C212/7143).
    ///
    /// Default: `false`. Only safe to enable on a PID-filtered MIG whose
    /// per-slot codes reflect the AHB-narrowed allowed sets — on the raw
    /// MIG, each slot's codes cover only one variant and strict matching
    /// would leave most segments unconsumed.
    pub strict_code_matching: bool,
}

/// MIG-guided assembler.
///
/// Takes a MIG schema and uses it as a grammar to guide consumption
/// of parsed EDIFACT segments. Produces a generic `AssembledTree`.
pub struct Assembler<'a> {
    mig: &'a MigSchema,
    config: AssemblerConfig,
}

impl<'a> Assembler<'a> {
    pub fn new(mig: &'a MigSchema) -> Self {
        Self {
            mig,
            config: AssemblerConfig::default(),
        }
    }

    pub fn with_config(mig: &'a MigSchema, config: AssemblerConfig) -> Self {
        Self { mig, config }
    }

    /// Assemble segments into a generic tree following MIG structure.
    pub fn assemble_generic(
        &self,
        segments: &[OwnedSegment],
    ) -> Result<AssembledTree, AssemblyError> {
        let mut cursor = SegmentCursor::new(segments.len());
        let mut tree = AssembledTree {
            segments: Vec::new(),
            groups: Vec::new(),
            post_group_start: 0,
            inter_group_segments: std::collections::BTreeMap::new(),
        };

        // Track which MIG segment indices were matched in the first pass
        let mut matched_seg_indices = Vec::new();

        // Process top-level segments (first pass — before groups)
        for (i, mig_seg) in self.mig.segments.iter().enumerate() {
            if cursor.is_exhausted() {
                break;
            }
            if let Some(assembled) = self.try_consume_segment(segments, &mut cursor, mig_seg)? {
                tree.segments.push(assembled);
                matched_seg_indices.push(i);
            }
        }

        // Process segment groups, interleaving root segment consumption.
        // Some message types (e.g., MSCONS) have root segments like UNS
        // between groups (SG2 and SG5). Before trying each group, consume
        // any unmatched root segments at the current cursor position.
        //
        // When consecutive same-ID groups have variant_code set (e.g., 3 SG8
        // entries for ZD7, Z98, ZF3), the assembler tries ALL variants at each
        // cursor position to handle interleaved reps.
        let mut group_idx = 0;
        while group_idx < self.mig.segment_groups.len() {
            if cursor.is_exhausted() {
                break;
            }

            let mig_group = &self.mig.segment_groups[group_idx];

            // Try consuming unmatched root segments before this group
            let tree_group_idx = tree.groups.len();
            for (i, mig_seg) in self.mig.segments.iter().enumerate() {
                if cursor.is_exhausted() {
                    break;
                }
                if matched_seg_indices.contains(&i) {
                    continue;
                }
                if let Some(assembled) = self.try_consume_segment(segments, &mut cursor, mig_seg)? {
                    tree.inter_group_segments
                        .entry(tree_group_idx)
                        .or_default()
                        .push(assembled);
                    matched_seg_indices.push(i);
                }
            }

            // Check if this starts a variant set (consecutive same-ID groups with variant_code)
            if mig_group.variant_code.is_some() {
                let variant_count = self.mig.segment_groups[group_idx..]
                    .iter()
                    .take_while(|g| g.id == mig_group.id && g.variant_code.is_some())
                    .count();
                let variant_end = group_idx + variant_count;

                let variant_groups = &self.mig.segment_groups[group_idx..variant_end];
                if let Some(combined) =
                    self.try_consume_variant_groups(segments, &mut cursor, variant_groups)?
                {
                    tree.groups.push(combined);
                }
                group_idx = variant_end;
            } else {
                if let Some(assembled) = self.try_consume_group(segments, &mut cursor, mig_group)? {
                    tree.groups.push(assembled);
                }
                group_idx += 1;
            }
        }

        // Mark where post-group segments start
        tree.post_group_start = tree.segments.len();

        // Second pass: try unmatched top-level segments (e.g., UNT, UNZ after groups)
        for (i, mig_seg) in self.mig.segments.iter().enumerate() {
            if cursor.is_exhausted() {
                break;
            }
            if matched_seg_indices.contains(&i) {
                continue;
            }
            if let Some(assembled) = self.try_consume_segment(segments, &mut cursor, mig_seg)? {
                tree.segments.push(assembled);
            }
        }

        Ok(tree)
    }

    fn try_consume_segment(
        &self,
        segments: &[OwnedSegment],
        cursor: &mut SegmentCursor,
        mig_seg: &MigSegment,
    ) -> Result<Option<AssembledSegment>, AssemblyError> {
        if cursor.is_exhausted() {
            return Ok(None);
        }
        let seg = &segments[cursor.position()];
        if matcher::matches_segment_tag(&seg.id, &mig_seg.id) {
            // Qualifier check: if the MIG slot has a qualifier_map entry,
            // verify the input segment's qualifier matches before consuming.
            if let Some(ref num) = mig_seg.number {
                if let Some((el_idx, comp_idx, expected)) = self.config.qualifier_map.get(num) {
                    let actual = seg
                        .elements
                        .get(*el_idx)
                        .and_then(|e| e.get(*comp_idx))
                        .map(|s| s.as_str())
                        .unwrap_or("");
                    if actual != expected {
                        return Ok(None); // Wrong qualifier — skip this slot
                    }
                }
            }
            // Note: full-code-profile matching (for disambiguating merged
            // sibling slots sharing a primary qualifier, e.g. PID 55035 PIA
            // variants all use 4347='5' but differ at C212/7143) is handled
            // by the caller in `try_consume_group`'s entry-run when strict
            // mode is on and `run_len > 1`. Doing it here would reject
            // solo-slot segments whose codes fall outside the AHB-narrowed
            // allowed set — those belong to the validator as COD002.
            let mut assembled = owned_to_assembled(seg);
            assembled.mig_number = mig_seg.number.clone();
            cursor.advance();
            Ok(Some(assembled))
        } else {
            Ok(None) // Segment not present (optional)
        }
    }

    /// Consume the entry run of a group with best-match slot selection.
    ///
    /// Used when `strict_code_matching` is on and `run_len > 1` — the group
    /// has multiple sibling entry slots with the same tag (merged PID-specific
    /// variants). For each pending segment, picks the unused slot whose full
    /// code profile matches best. Ties broken by MIG order. When no slot's
    /// profile matches, falls back to the first unused tag+qualifier-matching
    /// slot so the segment is still consumed (the validator emits COD002 if
    /// the code is truly invalid).
    fn consume_entry_run_best_match(
        &self,
        segments: &[OwnedSegment],
        cursor: &mut SegmentCursor,
        entry_slots: &[MigSegment],
        instance: &mut AssembledGroupInstance,
    ) -> Result<(), AssemblyError> {
        let mut used = vec![false; entry_slots.len()];
        for _ in 0..entry_slots.len() {
            if cursor.is_exhausted() {
                break;
            }
            let seg = &segments[cursor.position()];
            let mut strict_match: Option<usize> = None;
            let mut tag_match: Option<usize> = None;
            for (i, slot) in entry_slots.iter().enumerate() {
                if used[i] {
                    continue;
                }
                if !matcher::matches_segment_tag(&seg.id, &slot.id) {
                    continue;
                }
                if !self.segment_passes_qualifier_map(seg, slot) {
                    continue;
                }
                if tag_match.is_none() {
                    tag_match = Some(i);
                }
                if strict_match.is_none() && segment_matches_mig_codes(seg, slot) {
                    strict_match = Some(i);
                }
            }
            let Some(i) = strict_match.or(tag_match) else {
                break;
            };
            used[i] = true;
            let slot = &entry_slots[i];
            let mut assembled = owned_to_assembled(seg);
            assembled.mig_number = slot.number.clone();
            instance.segments.push(assembled);
            cursor.advance();
        }
        Ok(())
    }

    fn segment_passes_qualifier_map(&self, seg: &OwnedSegment, mig_seg: &MigSegment) -> bool {
        let Some(ref num) = mig_seg.number else {
            return true;
        };
        let Some((el_idx, comp_idx, expected)) = self.config.qualifier_map.get(num) else {
            return true;
        };
        let actual = seg
            .elements
            .get(*el_idx)
            .and_then(|e| e.get(*comp_idx))
            .map(|s| s.as_str())
            .unwrap_or("");
        actual == expected
    }

    fn try_consume_group(
        &self,
        segments: &[OwnedSegment],
        cursor: &mut SegmentCursor,
        mig_group: &MigSegmentGroup,
    ) -> Result<Option<AssembledGroup>, AssemblyError> {
        let mut repetitions = Vec::new();
        let entry_segment = mig_group.segments.first().ok_or_else(|| {
            AssemblyError::ParseError(format!("Group {} has no segments", mig_group.id))
        })?;

        // Loop for repeating groups
        while !cursor.is_exhausted() {
            let iter_start = cursor.position();
            let seg = &segments[cursor.position()];
            if !matcher::matches_segment_tag(&seg.id, &entry_segment.id) {
                break; // Current segment doesn't match group entry — stop repeating
            }

            // Check variant qualifier if set — tag matches but wrong variant
            if !mig_group.variant_codes.is_empty() {
                let (ei, ci) = mig_group.variant_qualifier_position.unwrap_or((0, 0));
                let actual_qual = seg
                    .elements
                    .get(ei)
                    .and_then(|e| e.get(ci))
                    .map(|s| s.as_str())
                    .unwrap_or("");
                if !mig_group
                    .variant_codes
                    .iter()
                    .any(|c| actual_qual.eq_ignore_ascii_case(c))
                {
                    break;
                }
            } else if let Some(ref expected_code) = mig_group.variant_code {
                let (ei, ci) = mig_group.variant_qualifier_position.unwrap_or((0, 0));
                let actual_qual = seg
                    .elements
                    .get(ei)
                    .and_then(|e| e.get(ci))
                    .map(|s| s.as_str())
                    .unwrap_or("");
                if !actual_qual.eq_ignore_ascii_case(expected_code) {
                    break;
                }
            }

            let mut instance = AssembledGroupInstance {
                segments: Vec::new(),
                child_groups: Vec::new(),
                entry_mig_number: entry_segment.number.clone(),
                variant_mig_numbers: collect_mig_numbers(mig_group),
                skipped_segments: Vec::new(),
            };

            // Consume segments within this group instance.
            // Process MIG slots in tag runs: for consecutive slots with the
            // same tag, consume ALL matching input segments — not just the
            // defined count. This handles real-world fixtures with more
            // repetitions than the merged MIG predicts (e.g., 6 RFFs when
            // the schema defines max 4).
            //
            // The entry segment (first tag run) is consumed bounded — one per
            // defined slot — because the outer while loop uses the entry tag
            // to delineate group repetitions.
            let mut slot_idx = 0;
            let mut is_entry_run = true;
            while slot_idx < mig_group.segments.len() {
                if cursor.is_exhausted() {
                    break;
                }
                let current_tag = &mig_group.segments[slot_idx].id;
                let run_len = mig_group.segments[slot_idx..]
                    .iter()
                    .take_while(|s| s.id == *current_tag)
                    .count();

                if is_entry_run {
                    // Entry tag: consume at most run_len (preserves group boundaries)
                    let entry_slots = &mig_group.segments[slot_idx..slot_idx + run_len];
                    if self.config.strict_code_matching && run_len > 1 {
                        // Best-match: among tag-matching sibling slots, prefer
                        // the one whose full code profile matches the segment
                        // (disambiguates PID 55035 PIA 00108/Z12 vs 00197/SRW).
                        // Falls back to MIG order when no profile matches —
                        // preserves assembly for codes outside any AHB-narrowed set.
                        self.consume_entry_run_best_match(
                            segments,
                            cursor,
                            entry_slots,
                            &mut instance,
                        )?;
                    } else {
                        for slot in entry_slots {
                            if cursor.is_exhausted() {
                                break;
                            }
                            if let Some(assembled) =
                                self.try_consume_segment(segments, cursor, slot)?
                            {
                                instance.segments.push(assembled);
                            }
                        }
                    }
                    is_entry_run = false;
                } else if matcher::matches_segment_tag(current_tag, &entry_segment.id) {
                    // Non-entry slot with SAME tag as entry (e.g., CCI appears as
                    // both entry and non-entry in merged SG30).
                    //
                    // Only consume if we haven't yet consumed any NON-entry-tag
                    // segments (i.e., we're still in a consecutive entry-tag run).
                    // Once we've consumed a different tag (like CAV), seeing the
                    // entry tag again means a new rep boundary.
                    //
                    // z35: entry CCI → CAV CAV → sees CCI → has_other=true → break ✓
                    // z39: entry CCI → (no CAV) → sees CCI → has_other=false → consume ✓
                    //      then CCI CCI → CAV → sees CCI → has_other=true → break
                    //      BUT: z39 needs CCI-CAV-CCI-CAV structure
                    //
                    // Better heuristic: check if ALL remaining slots from here are
                    // entry-tag + non-entry pairs. If the current slot is entry-tag
                    // and the NEXT input segment after it would be a non-entry tag,
                    // consume — it's a continuation. Otherwise break.
                    if cursor.is_exhausted() {
                        break;
                    }
                    let seg = &segments[cursor.position()];
                    if !matcher::matches_segment_tag(&seg.id, current_tag) {
                        break;
                    }
                    // Check: is there a non-entry segment AFTER this entry-tag?
                    // If so, this CCI+CAV pair is part of the current rep.
                    let has_following_non_entry = if cursor.position() + 1 < segments.len() {
                        let next = &segments[cursor.position() + 1];
                        !matcher::matches_segment_tag(&next.id, &entry_segment.id)
                            && mig_group.segments.iter().any(|s| {
                                matcher::matches_segment_tag(&next.id, &s.id)
                                    && !matcher::matches_segment_tag(&s.id, &entry_segment.id)
                            })
                    } else {
                        false
                    };
                    if has_following_non_entry {
                        // CCI followed by CAV → consume as continuation pair
                        instance.segments.push(owned_to_assembled(seg));
                        cursor.advance();
                    } else {
                        // CCI followed by CCI or unknown → let outer loop decide
                        break;
                    }
                } else {
                    // Non-entry tag: consume bounded slots first (with mig_number),
                    // then greedily consume extras (without mig_number).
                    // The bounded slots get mig_number from the MIG definition so
                    // the validator can distinguish same-tag segments (e.g., DTM+92
                    // vs DTM+93 both in SG4).
                    let slots = &mig_group.segments[slot_idx..slot_idx + run_len];
                    if self.config.strict_code_matching && run_len > 1 {
                        self.consume_entry_run_best_match(segments, cursor, slots, &mut instance)?;
                    } else {
                        for slot in slots {
                            if cursor.is_exhausted() {
                                break;
                            }
                            if let Some(assembled) =
                                self.try_consume_segment(segments, cursor, slot)?
                            {
                                instance.segments.push(assembled);
                            }
                        }
                    }
                    // Greedily consume any remaining same-tag segments beyond the MIG count
                    while !cursor.is_exhausted() {
                        let seg = &segments[cursor.position()];
                        if matcher::matches_segment_tag(&seg.id, current_tag) {
                            instance.segments.push(owned_to_assembled(seg));
                            cursor.advance();
                        } else {
                            break;
                        }
                    }
                }

                slot_idx += run_len;

                // Point A: Skip unknown segments between MIG slot runs.
                // When skip mode is ON and we just finished a slot run but the
                // current segment doesn't match any remaining MIG slot, nested
                // group entry, or the entry tag, skip it.
                if self.config.skip_unknown_segments {
                    while !cursor.is_exhausted() {
                        let seg = &segments[cursor.position()];
                        // Stop if it matches the entry tag (next group repetition)
                        if matcher::matches_segment_tag(&seg.id, &entry_segment.id) {
                            break;
                        }
                        // Stop if it matches any remaining MIG slot
                        if mig_group.segments[slot_idx..]
                            .iter()
                            .any(|s| matcher::matches_segment_tag(&seg.id, &s.id))
                        {
                            break;
                        }
                        // Stop if it matches any nested group entry
                        if mig_group.nested_groups.iter().any(|ng| {
                            ng.segments
                                .first()
                                .is_some_and(|es| matcher::matches_segment_tag(&seg.id, &es.id))
                        }) {
                            break;
                        }
                        // Unknown segment — skip it
                        instance.skipped_segments.push(owned_to_assembled(seg));
                        cursor.advance();
                    }
                }
            }

            // Consume nested groups (variant-aware for same-ID groups)
            let mut nested_idx = 0;
            while nested_idx < mig_group.nested_groups.len() {
                if cursor.is_exhausted() {
                    break;
                }
                let nested = &mig_group.nested_groups[nested_idx];

                if nested.variant_code.is_some() {
                    // Variant set: collect consecutive same-ID groups with variant_code
                    let variant_count = mig_group.nested_groups[nested_idx..]
                        .iter()
                        .take_while(|g| g.id == nested.id && g.variant_code.is_some())
                        .count();
                    let variant_end = nested_idx + variant_count;
                    let variant_groups = &mig_group.nested_groups[nested_idx..variant_end];
                    if let Some(combined) =
                        self.try_consume_variant_groups(segments, cursor, variant_groups)?
                    {
                        instance.child_groups.push(combined);
                    }
                    nested_idx = variant_end;
                } else {
                    if let Some(assembled) = self.try_consume_group(segments, cursor, nested)? {
                        instance.child_groups.push(assembled);
                    }
                    nested_idx += 1;
                }
            }

            // Guard against infinite loops: if no progress was made this iteration
            // (entry tag matched but the entry segment was rejected by e.g. a
            // qualifier_map mismatch), stop. Pushing an empty rep per iteration
            // would allocate unbounded memory (see collect_mig_numbers call in
            // the instance constructor).
            if cursor.position() == iter_start {
                break;
            }
            repetitions.push(instance);
        }

        if repetitions.is_empty() {
            Ok(None)
        } else {
            Ok(Some(AssembledGroup {
                group_id: mig_group.id.clone(),
                repetitions,
            }))
        }
    }

    /// Consume interleaved repetitions of variant groups.
    ///
    /// At each cursor position, tries all variant definitions to find which one
    /// matches the entry segment's qualifier. Collects all reps into one
    /// `AssembledGroup` with the shared group_id.
    fn try_consume_variant_groups(
        &self,
        segments: &[OwnedSegment],
        cursor: &mut SegmentCursor,
        variants: &[MigSegmentGroup],
    ) -> Result<Option<AssembledGroup>, AssemblyError> {
        let group_id = variants[0].id.clone();
        let entry_tag = variants[0]
            .segments
            .first()
            .map(|s| s.id.as_str())
            .unwrap_or("");
        let mut all_reps = Vec::new();

        while !cursor.is_exhausted() {
            let seg = &segments[cursor.position()];
            if !matcher::matches_segment_tag(&seg.id, entry_tag) {
                break;
            }

            // Find which variant matches this segment's qualifier.
            // Each variant may have its qualifier at a different element position
            // (e.g., CCI+Z19 has qualifier at [0][0], but CCI+++Z15 at [2][0]).
            let matched = variants.iter().find(|v| {
                let (ei, ci) = v.variant_qualifier_position.unwrap_or((0, 0));
                let actual_qual = seg
                    .elements
                    .get(ei)
                    .and_then(|e| e.get(ci))
                    .map(|s| s.as_str())
                    .unwrap_or("");
                if !v.variant_codes.is_empty() {
                    v.variant_codes
                        .iter()
                        .any(|c| actual_qual.eq_ignore_ascii_case(c))
                } else if let Some(ref expected_code) = v.variant_code {
                    actual_qual.eq_ignore_ascii_case(expected_code)
                } else {
                    false
                }
            });

            if let Some(variant) = matched {
                if let Some(group) = self.try_consume_group(segments, cursor, variant)? {
                    all_reps.extend(group.repetitions);
                } else {
                    break;
                }
            } else {
                // No variant matches — try consuming with the first variant as
                // fallback to avoid getting stuck. This handles edge cases where
                // the qualifier doesn't exactly match any variant code.
                if let Some(group) = self.try_consume_group(segments, cursor, &variants[0])? {
                    all_reps.extend(group.repetitions);
                } else {
                    break;
                }
            }
        }

        if all_reps.is_empty() {
            Ok(None)
        } else {
            Ok(Some(AssembledGroup {
                group_id,
                repetitions: all_reps,
            }))
        }
    }

    /// Assemble segments with diagnostic collection.
    ///
    /// Returns the assembled tree plus diagnostics for segments not consumed
    /// by the MIG-guided assembly. Existing `assemble_generic()` is unchanged.
    pub fn assemble_with_diagnostics(
        &self,
        segments: &[OwnedSegment],
    ) -> (AssembledTree, Vec<StructureDiagnostic>) {
        let mut diagnostics = Vec::new();

        let tree = match self.assemble_generic(segments) {
            Ok(tree) => tree,
            Err(e) => {
                diagnostics.push(StructureDiagnostic {
                    kind: StructureDiagnosticKind::UnexpectedSegment,
                    segment_id: String::new(),
                    position: 0,
                    message: format!("Assembly failed: {e}"),
                });
                return (
                    AssembledTree {
                        segments: Vec::new(),
                        groups: Vec::new(),
                        post_group_start: 0,
                        inter_group_segments: std::collections::BTreeMap::new(),
                    },
                    diagnostics,
                );
            }
        };

        // Count consumed segments in the assembled tree
        let consumed = count_tree_segments(&tree);

        // Segments beyond consumed count are unconsumed
        for (i, seg) in segments.iter().enumerate().skip(consumed) {
            diagnostics.push(StructureDiagnostic {
                kind: StructureDiagnosticKind::UnexpectedSegment,
                segment_id: seg.id.clone(),
                position: i,
                message: format!(
                    "Segment '{}' at position {} was not consumed by MIG-guided assembly",
                    seg.id, i
                ),
            });
        }

        (tree, diagnostics)
    }
}

fn count_tree_segments(tree: &AssembledTree) -> usize {
    let mut count = tree.segments.len();
    for group in &tree.groups {
        count += count_group_segments(group);
    }
    // Count inter-group segments (e.g., UNS+D between groups)
    for segs in tree.inter_group_segments.values() {
        count += segs.len();
    }
    count
}

fn count_group_segments(group: &AssembledGroup) -> usize {
    let mut count = 0;
    for rep in &group.repetitions {
        count += rep.segments.len();
        count += rep.skipped_segments.len();
        for child in &rep.child_groups {
            count += count_group_segments(child);
        }
    }
    count
}

/// Collect all MIG `Number`s from a segment group definition, recursively.
///
/// This includes numbers from direct segments and from nested groups.
/// Used to populate `AssembledGroupInstance::variant_mig_numbers`.
fn collect_mig_numbers(group: &MigSegmentGroup) -> Vec<String> {
    let mut numbers = Vec::new();
    for seg in &group.segments {
        if let Some(ref num) = seg.number {
            numbers.push(num.clone());
        }
    }
    for nested in &group.nested_groups {
        numbers.extend(collect_mig_numbers(nested));
    }
    numbers
}

pub fn owned_to_assembled(seg: &OwnedSegment) -> AssembledSegment {
    AssembledSegment {
        tag: seg.id.clone(),
        elements: seg.elements.clone(),
        mig_number: None,
    }
}

/// Check every code-bearing position declared on a MIG segment against the
/// corresponding value on the input segment.
///
/// Used by `try_consume_segment` to disambiguate slots that share the same
/// primary qualifier but differ on a secondary code. Returns `true` when the
/// input segment's values at each declared position are either empty
/// (optional) or in the slot's allowed set.
fn segment_matches_mig_codes(seg: &OwnedSegment, mig_seg: &MigSegment) -> bool {
    let actual_at = |el: usize, c: usize| -> &str {
        seg.elements
            .get(el)
            .and_then(|e| e.get(c))
            .map(|s| s.as_str())
            .unwrap_or("")
    };
    for de in &mig_seg.data_elements {
        if !de.codes.is_empty() {
            let actual = actual_at(de.position, 0);
            if !actual.is_empty() && !de.codes.iter().any(|c| c.value == actual) {
                return false;
            }
        }
    }
    for comp in &mig_seg.composites {
        for de in &comp.data_elements {
            if !de.codes.is_empty() {
                let actual = actual_at(comp.position, de.position);
                if !actual.is_empty() && !de.codes.iter().any(|c| c.value == actual) {
                    return false;
                }
            }
        }
    }
    true
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::test_support::{make_mig_group, make_mig_group_with_variant, make_mig_segment};

    fn make_owned_seg(id: &str, elements: Vec<Vec<&str>>) -> OwnedSegment {
        OwnedSegment {
            id: id.to_string(),
            elements: elements
                .into_iter()
                .map(|e| e.into_iter().map(|c| c.to_string()).collect())
                .collect(),
            segment_number: 0,
        }
    }

    fn make_mig_schema(segments: Vec<&str>, groups: Vec<MigSegmentGroup>) -> MigSchema {
        MigSchema {
            message_type: "UTILMD".to_string(),
            variant: Some("Strom".to_string()),
            version: "S2.1".to_string(),
            publication_date: "2025-03-20".to_string(),
            author: "BDEW".to_string(),
            format_version: "FV2504".to_string(),
            source_file: "test".to_string(),
            segments: segments.into_iter().map(make_mig_segment).collect(),
            segment_groups: groups,
        }
    }

    #[test]
    fn test_assembler_top_level_segments_only() {
        let mig = make_mig_schema(vec!["UNH", "BGM", "DTM", "UNT"], vec![]);

        let segments = vec![
            make_owned_seg("UNH", vec![vec!["001", "UTILMD:D:11A:UN:S2.1"]]),
            make_owned_seg("BGM", vec![vec!["E01", "DOC001"]]),
            make_owned_seg("DTM", vec![vec!["137", "20250101", "102"]]),
            make_owned_seg("UNT", vec![vec!["4", "001"]]),
        ];

        let assembler = Assembler::new(&mig);
        let result = assembler.assemble_generic(&segments).unwrap();

        assert_eq!(result.segments.len(), 4);
        assert_eq!(result.segments[0].tag, "UNH");
        assert_eq!(result.segments[1].tag, "BGM");
        assert_eq!(result.segments[2].tag, "DTM");
        assert_eq!(result.segments[3].tag, "UNT");
        assert!(result.groups.is_empty());
    }

    #[test]
    fn test_assembler_with_segment_group() {
        let mig = make_mig_schema(
            vec!["UNH", "BGM"],
            vec![
                make_mig_group("SG2", vec!["NAD"], vec![]),
                make_mig_group("SG4", vec!["IDE", "STS"], vec![]),
            ],
        );

        let segments = vec![
            make_owned_seg("UNH", vec![vec!["001"]]),
            make_owned_seg("BGM", vec![vec!["E01"]]),
            make_owned_seg("NAD", vec![vec!["MS", "9900123"]]),
            make_owned_seg("NAD", vec![vec!["MR", "9900456"]]),
            make_owned_seg("IDE", vec![vec!["24", "TX001"]]),
            make_owned_seg("STS", vec![vec!["7"], vec!["Z33"]]),
        ];

        let assembler = Assembler::new(&mig);
        let result = assembler.assemble_generic(&segments).unwrap();

        // Top-level: UNH, BGM
        assert_eq!(result.segments.len(), 2);
        // SG2: 2 repetitions (two NAD segments)
        assert_eq!(result.groups.len(), 2);
        assert_eq!(result.groups[0].group_id, "SG2");
        assert_eq!(result.groups[0].repetitions.len(), 2);
        assert_eq!(result.groups[0].repetitions[0].segments[0].tag, "NAD");
        assert_eq!(result.groups[0].repetitions[1].segments[0].tag, "NAD");
        // SG4: 1 repetition (IDE + STS)
        assert_eq!(result.groups[1].group_id, "SG4");
        assert_eq!(result.groups[1].repetitions.len(), 1);
        assert_eq!(result.groups[1].repetitions[0].segments.len(), 2);
    }

    #[test]
    fn test_assembler_nested_groups() {
        let sg3 = make_mig_group("SG3", vec!["CTA", "COM"], vec![]);
        let mig = make_mig_schema(
            vec!["UNH", "BGM"],
            vec![make_mig_group("SG2", vec!["NAD"], vec![sg3])],
        );

        let segments = vec![
            make_owned_seg("UNH", vec![vec!["001"]]),
            make_owned_seg("BGM", vec![vec!["E01"]]),
            make_owned_seg("NAD", vec![vec!["MS", "9900123"]]),
            make_owned_seg("CTA", vec![vec!["IC", "Kontakt"]]),
            make_owned_seg("COM", vec![vec!["040@example.com", "EM"]]),
        ];

        let assembler = Assembler::new(&mig);
        let result = assembler.assemble_generic(&segments).unwrap();

        // SG2 has 1 repetition
        let sg2 = &result.groups[0];
        assert_eq!(sg2.group_id, "SG2");
        assert_eq!(sg2.repetitions.len(), 1);

        let sg2_inst = &sg2.repetitions[0];
        assert_eq!(sg2_inst.segments[0].tag, "NAD");

        // SG3 nested inside SG2
        assert_eq!(sg2_inst.child_groups.len(), 1);
        let sg3 = &sg2_inst.child_groups[0];
        assert_eq!(sg3.group_id, "SG3");
        assert_eq!(sg3.repetitions[0].segments.len(), 2);
        assert_eq!(sg3.repetitions[0].segments[0].tag, "CTA");
        assert_eq!(sg3.repetitions[0].segments[1].tag, "COM");
    }

    #[test]
    fn test_assembler_optional_segments_skipped() {
        // MIG expects UNH, BGM, DTM, UNT but input has no DTM
        let mig = make_mig_schema(vec!["UNH", "BGM", "DTM", "UNT"], vec![]);

        let segments = vec![
            make_owned_seg("UNH", vec![vec!["001"]]),
            make_owned_seg("BGM", vec![vec!["E01"]]),
            make_owned_seg("UNT", vec![vec!["2", "001"]]),
        ];

        let assembler = Assembler::new(&mig);
        let result = assembler.assemble_generic(&segments).unwrap();

        // DTM is skipped (optional), UNT consumed
        assert_eq!(result.segments.len(), 3);
        assert_eq!(result.segments[0].tag, "UNH");
        assert_eq!(result.segments[1].tag, "BGM");
        assert_eq!(result.segments[2].tag, "UNT");
    }

    #[test]
    fn test_assembler_empty_segments() {
        let mig = make_mig_schema(vec!["UNH"], vec![]);
        let assembler = Assembler::new(&mig);
        let result = assembler.assemble_generic(&[]).unwrap();
        assert!(result.segments.is_empty());
        assert!(result.groups.is_empty());
    }

    #[test]
    fn test_assembler_preserves_element_data() {
        let mig = make_mig_schema(vec!["DTM"], vec![]);

        let segments = vec![make_owned_seg(
            "DTM",
            vec![vec!["137", "202501010000+01", "303"]],
        )];

        let assembler = Assembler::new(&mig);
        let result = assembler.assemble_generic(&segments).unwrap();

        let dtm = &result.segments[0];
        assert_eq!(dtm.elements[0][0], "137");
        assert_eq!(dtm.elements[0][1], "202501010000+01");
        assert_eq!(dtm.elements[0][2], "303");
    }

    #[test]
    fn test_group_instance_as_assembled_tree() {
        // Build an SG4 instance with root segments (IDE, STS) and child groups (SG5)
        let sg5 = AssembledGroup {
            group_id: "SG5".to_string(),
            repetitions: vec![AssembledGroupInstance {
                segments: vec![AssembledSegment {
                    tag: "LOC".to_string(),
                    elements: vec![vec!["Z16".to_string(), "DE000111222333".to_string()]],
                    mig_number: None,
                }],
                child_groups: vec![],
                entry_mig_number: None,
                variant_mig_numbers: vec![],
                skipped_segments: vec![],
            }],
        };

        let sg4_instance = AssembledGroupInstance {
            segments: vec![
                AssembledSegment {
                    tag: "IDE".to_string(),
                    elements: vec![vec!["24".to_string(), "TX001".to_string()]],
                    mig_number: None,
                },
                AssembledSegment {
                    tag: "STS".to_string(),
                    elements: vec![vec!["7".to_string()]],
                    mig_number: None,
                },
            ],
            child_groups: vec![sg5],
            entry_mig_number: None,
            variant_mig_numbers: vec![],
            skipped_segments: vec![],
        };

        let sub_tree = sg4_instance.as_assembled_tree();

        // Root segments of sub-tree are the SG4 instance's segments
        assert_eq!(sub_tree.segments.len(), 2);
        assert_eq!(sub_tree.segments[0].tag, "IDE");
        assert_eq!(sub_tree.segments[1].tag, "STS");

        // Groups of sub-tree are the SG4 instance's child groups
        assert_eq!(sub_tree.groups.len(), 1);
        assert_eq!(sub_tree.groups[0].group_id, "SG5");

        // post_group_start marks where root segments end
        assert_eq!(sub_tree.post_group_start, 2);
    }

    #[test]
    fn test_assembler_from_parsed_edifact() {
        // End-to-end: parse raw EDIFACT, then assemble
        let input = b"UNA:+.? 'UNB+UNOC:3+SENDER+RECEIVER+210101:1200+REF001'UNH+MSG001+UTILMD:D:11A:UN:S2.1'BGM+E01+DOC001+9'DTM+137:20250101:102'UNT+3+MSG001'UNZ+1+REF001'";
        let segments = crate::tokenize::parse_to_segments(input).unwrap();

        let mig = make_mig_schema(vec!["UNB", "UNH", "BGM", "DTM", "UNT", "UNZ"], vec![]);

        let assembler = Assembler::new(&mig);
        let result = assembler.assemble_generic(&segments).unwrap();

        assert!(result.segments.iter().any(|s| s.tag == "UNH"));
        assert!(result.segments.iter().any(|s| s.tag == "BGM"));
        assert!(result.segments.iter().any(|s| s.tag == "DTM"));
    }

    #[test]
    fn test_assemble_with_diagnostics_clean_input() {
        let mig = make_mig_schema(vec!["UNH", "BGM", "UNT"], vec![]);
        let segments = vec![
            make_owned_seg("UNH", vec![vec!["001"]]),
            make_owned_seg("BGM", vec![vec!["E01"]]),
            make_owned_seg("UNT", vec![vec!["2", "001"]]),
        ];
        let assembler = Assembler::new(&mig);
        let (tree, diagnostics) = assembler.assemble_with_diagnostics(&segments);
        assert_eq!(tree.segments.len(), 3);
        assert!(
            diagnostics.is_empty(),
            "Clean input should have no diagnostics"
        );
    }

    #[test]
    fn test_assemble_with_diagnostics_unconsumed_segments() {
        let mig = make_mig_schema(vec!["UNH", "BGM"], vec![]);
        let segments = vec![
            make_owned_seg("UNH", vec![vec!["001"]]),
            make_owned_seg("BGM", vec![vec!["E01"]]),
            make_owned_seg("FTX", vec![vec!["AAA", "extra text"]]),
        ];
        let assembler = Assembler::new(&mig);
        let (tree, diagnostics) = assembler.assemble_with_diagnostics(&segments);
        assert_eq!(tree.segments.len(), 2);
        assert_eq!(diagnostics.len(), 1);
        assert_eq!(
            diagnostics[0].kind,
            StructureDiagnosticKind::UnexpectedSegment
        );
        assert_eq!(diagnostics[0].segment_id, "FTX");
        assert_eq!(diagnostics[0].position, 2);
    }

    #[test]
    fn test_assemble_with_diagnostics_multiple_unconsumed() {
        let mig = make_mig_schema(vec!["UNH"], vec![]);
        let segments = vec![
            make_owned_seg("UNH", vec![vec!["001"]]),
            make_owned_seg("FOO", vec![]),
            make_owned_seg("BAR", vec![]),
            make_owned_seg("BAZ", vec![]),
        ];
        let assembler = Assembler::new(&mig);
        let (tree, diagnostics) = assembler.assemble_with_diagnostics(&segments);
        assert_eq!(tree.segments.len(), 1);
        assert_eq!(diagnostics.len(), 3);
        assert_eq!(diagnostics[0].segment_id, "FOO");
        assert_eq!(diagnostics[1].segment_id, "BAR");
        assert_eq!(diagnostics[2].segment_id, "BAZ");
    }

    // ── Non-entry segment mig_number assignment tests ──

    #[test]
    fn test_non_entry_segments_get_mig_number_from_bounded_slots() {
        // MIG group SG4 has entry IDE + two numbered DTMs + STS.
        // The assembler should assign mig_number from the MIG slots to
        // each non-entry segment via the bounded consumption path.
        use crate::test_support::make_mig_segment_numbered;

        let sg4 = MigSegmentGroup {
            segments: vec![
                make_mig_segment_numbered("IDE", "00020"),
                make_mig_segment_numbered("DTM", "00023"),
                make_mig_segment_numbered("DTM", "00024"),
                make_mig_segment_numbered("STS", "00035"),
            ],
            ..make_mig_group("SG4", vec![], vec![])
        };
        let mig = make_mig_schema(vec!["UNH"], vec![sg4]);

        let segments = vec![
            make_owned_seg("UNH", vec![vec!["001"]]),
            make_owned_seg("IDE", vec![vec!["24", "TX001"]]),
            make_owned_seg("DTM", vec![vec!["92", "202505312200+00", "303"]]),
            make_owned_seg("DTM", vec![vec!["93", "202512312300+00", "303"]]),
            make_owned_seg("STS", vec![vec!["7"], vec![], vec!["E01"]]),
        ];

        let assembler = Assembler::new(&mig);
        let tree = assembler.assemble_generic(&segments).unwrap();

        let sg4_instance = &tree.groups[0].repetitions[0];

        // IDE (entry) gets mig_number from try_consume_segment
        assert_eq!(sg4_instance.segments[0].tag, "IDE");
        assert_eq!(sg4_instance.segments[0].mig_number.as_deref(), Some("00020"));

        // DTM+92 gets mig_number "00023" from first DTM slot
        assert_eq!(sg4_instance.segments[1].tag, "DTM");
        assert_eq!(sg4_instance.segments[1].mig_number.as_deref(), Some("00023"));

        // DTM+93 gets mig_number "00024" from second DTM slot
        assert_eq!(sg4_instance.segments[2].tag, "DTM");
        assert_eq!(sg4_instance.segments[2].mig_number.as_deref(), Some("00024"));

        // STS gets mig_number "00035"
        assert_eq!(sg4_instance.segments[3].tag, "STS");
        assert_eq!(sg4_instance.segments[3].mig_number.as_deref(), Some("00035"));

        // variant_mig_numbers should contain all four
        assert!(sg4_instance.variant_mig_numbers.contains(&"00020".to_string()));
        assert!(sg4_instance.variant_mig_numbers.contains(&"00023".to_string()));
        assert!(sg4_instance.variant_mig_numbers.contains(&"00024".to_string()));
        assert!(sg4_instance.variant_mig_numbers.contains(&"00035".to_string()));
    }

    #[test]
    fn test_greedy_extra_segments_get_no_mig_number() {
        // MIG defines 1 DTM slot, but input has 2 DTMs.
        // First DTM gets mig_number from bounded path, second gets None (greedy extra).
        use crate::test_support::make_mig_segment_numbered;

        let sg4 = MigSegmentGroup {
            segments: vec![
                make_mig_segment_numbered("IDE", "00020"),
                make_mig_segment_numbered("DTM", "00023"),
            ],
            ..make_mig_group("SG4", vec![], vec![])
        };
        let mig = make_mig_schema(vec!["UNH"], vec![sg4]);

        let segments = vec![
            make_owned_seg("UNH", vec![vec!["001"]]),
            make_owned_seg("IDE", vec![vec!["24"]]),
            make_owned_seg("DTM", vec![vec!["92", "20250531"]]),
            make_owned_seg("DTM", vec![vec!["93", "20251231"]]), // extra beyond MIG
        ];

        let assembler = Assembler::new(&mig);
        let tree = assembler.assemble_generic(&segments).unwrap();

        let sg4_instance = &tree.groups[0].repetitions[0];
        assert_eq!(sg4_instance.segments.len(), 3); // IDE + 2 DTMs

        // First DTM: bounded slot → mig_number set
        assert_eq!(sg4_instance.segments[1].mig_number.as_deref(), Some("00023"));

        // Second DTM: greedy extra → mig_number None
        assert_eq!(sg4_instance.segments[2].mig_number, None);
    }

    // ── Qualifier-aware assembly tests ──

    #[test]
    fn test_qualifier_map_prevents_wrong_slot_consumption() {
        // MIG defines DTM(00023) + DTM(00024). Input has only DTM+93.
        // Without qualifier map: DTM+93 consumed by slot 00023 (wrong).
        // With qualifier map: slot 00023 expects "92", skips DTM+93.
        //   Slot 00024 expects "93", consumes DTM+93 correctly.
        use crate::test_support::make_mig_segment_numbered;
        use std::collections::HashMap;

        let sg4 = MigSegmentGroup {
            segments: vec![
                make_mig_segment_numbered("IDE", "00020"),
                make_mig_segment_numbered("DTM", "00023"),
                make_mig_segment_numbered("DTM", "00024"),
            ],
            ..make_mig_group("SG4", vec![], vec![])
        };
        let mig = make_mig_schema(vec!["UNH"], vec![sg4]);

        let segments = vec![
            make_owned_seg("UNH", vec![vec!["001"]]),
            make_owned_seg("IDE", vec![vec!["24"]]),
            make_owned_seg("DTM", vec![vec!["93", "202512312300+00", "303"]]),
        ];

        let mut qualifier_map = HashMap::new();
        qualifier_map.insert("00023".to_string(), (0, 0, "92".to_string()));
        qualifier_map.insert("00024".to_string(), (0, 0, "93".to_string()));

        let config = AssemblerConfig {
            skip_unknown_segments: false,
            qualifier_map,
            ..Default::default()
        };
        let assembler = Assembler::with_config(&mig, config);
        let tree = assembler.assemble_generic(&segments).unwrap();

        let sg4_instance = &tree.groups[0].repetitions[0];

        // DTM+93 should be consumed by slot 00024, NOT slot 00023
        assert_eq!(sg4_instance.segments.len(), 2); // IDE + DTM+93
        let dtm = &sg4_instance.segments[1];
        assert_eq!(dtm.tag, "DTM");
        assert_eq!(
            dtm.mig_number.as_deref(),
            Some("00024"),
            "DTM+93 should get mig_number 00024 (not 00023)"
        );
    }

    #[test]
    fn test_group_entry_qualifier_mismatch_does_not_infinite_loop() {
        // Regression: when a group's entry segment has a qualifier_map entry
        // but the input segment's qualifier does not match, the outer
        // `while !cursor.is_exhausted()` loop in try_consume_group used to
        // spin forever — entry tag matched, so the loop kept going, but
        // try_consume_segment rejected the segment on qualifier mismatch, so
        // the cursor never advanced. Each iteration allocated a fresh
        // variant_mig_numbers Vec via collect_mig_numbers, driving unbounded
        // memory growth (observed: 3.5 GB → OOM in ~3 s on staging for
        // FV2604/UTILMD_Gas/PID 44004 with LOC+172).
        use crate::test_support::make_mig_segment_numbered;
        use std::collections::HashMap;

        let sg5 = MigSegmentGroup {
            segments: vec![make_mig_segment_numbered("LOC", "00050")],
            ..make_mig_group("SG5", vec![], vec![])
        };
        let mig = make_mig_schema(vec!["UNH"], vec![sg5]);

        let segments = vec![
            make_owned_seg("UNH", vec![vec!["001"]]),
            // LOC entry tag matches, but qualifier "172" ≠ expected "Z16"
            make_owned_seg("LOC", vec![vec!["172"], vec!["92003964705"]]),
        ];

        let mut qualifier_map = HashMap::new();
        qualifier_map.insert("00050".to_string(), (0, 0, "Z16".to_string()));

        let config = AssemblerConfig {
            skip_unknown_segments: false,
            qualifier_map,
            ..Default::default()
        };
        let assembler = Assembler::with_config(&mig, config);

        // Before the fix this would loop forever. Bound the assertion with a
        // generous wall-clock guard so a regression is a clear test failure
        // rather than a hanging CI job.
        let start = std::time::Instant::now();
        let tree = assembler.assemble_generic(&segments).unwrap();
        assert!(
            start.elapsed() < std::time::Duration::from_secs(5),
            "assembly took {:?} — suspected infinite-loop regression",
            start.elapsed()
        );

        // LOC+172 didn't match SG5's qualifier, so SG5 should be empty.
        // The LOC segment remains unconsumed (caller will surface it as a
        // structure diagnostic).
        assert!(tree.groups.is_empty());
    }

    // ── Skip-unknown-segments tests ──

    #[test]
    fn test_skip_unknown_segment_between_slots() {
        // MIG group expects [SEQ, CCI], input has [SEQ, RFF, CCI].
        // With skip ON, RFF is skipped and CCI is consumed.
        // With skip OFF (default), CCI is lost because RFF stalls the cursor.
        let sg8 = make_mig_group("SG8", vec!["SEQ", "CCI"], vec![]);
        let mig = make_mig_schema(vec!["UNH"], vec![sg8.clone()]);

        let segments = vec![
            make_owned_seg("UNH", vec![vec!["001"]]),
            make_owned_seg("SEQ", vec![vec!["Z98"]]),
            make_owned_seg("RFF", vec![vec!["Z38", "CROSSREF"]]),
            make_owned_seg("CCI", vec![vec!["Z30"]]),
        ];

        // Skip OFF: CCI not consumed (RFF stalls cursor after SEQ)
        let off = Assembler::new(&mig);
        let tree_off = off.assemble_generic(&segments).unwrap();
        let sg8_off = &tree_off.groups[0];
        assert_eq!(sg8_off.repetitions[0].segments.len(), 1); // Only SEQ
        assert_eq!(sg8_off.repetitions[0].segments[0].tag, "SEQ");

        // Skip ON: RFF skipped, CCI consumed
        let on = Assembler::with_config(
            &mig,
            AssemblerConfig {
                skip_unknown_segments: true,
                ..Default::default()
            },
        );
        let tree_on = on.assemble_generic(&segments).unwrap();
        let sg8_on = &tree_on.groups[0];
        assert_eq!(sg8_on.repetitions[0].segments.len(), 2); // SEQ + CCI
        assert_eq!(sg8_on.repetitions[0].segments[0].tag, "SEQ");
        assert_eq!(sg8_on.repetitions[0].segments[1].tag, "CCI");
    }

    #[test]
    fn test_skip_preserves_on_instance() {
        // Skipped segments are stored in instance.skipped_segments
        let sg8 = make_mig_group("SG8", vec!["SEQ", "CCI"], vec![]);
        let mig = make_mig_schema(vec!["UNH"], vec![sg8]);

        let segments = vec![
            make_owned_seg("UNH", vec![vec!["001"]]),
            make_owned_seg("SEQ", vec![vec!["Z98"]]),
            make_owned_seg("RFF", vec![vec!["Z38", "REF1"]]),
            make_owned_seg("DTM", vec![vec!["92", "20250101"]]),
            make_owned_seg("CCI", vec![vec!["Z30"]]),
        ];

        let assembler = Assembler::with_config(
            &mig,
            AssemblerConfig {
                skip_unknown_segments: true,
                ..Default::default()
            },
        );
        let tree = assembler.assemble_generic(&segments).unwrap();
        let instance = &tree.groups[0].repetitions[0];

        assert_eq!(instance.segments.len(), 2); // SEQ + CCI
        assert_eq!(instance.skipped_segments.len(), 2); // RFF + DTM
        assert_eq!(instance.skipped_segments[0].tag, "RFF");
        assert_eq!(instance.skipped_segments[1].tag, "DTM");
    }

    #[test]
    fn test_skip_mode_off_default() {
        // Assembler::new() doesn't skip (backwards compat)
        let mig = make_mig_schema(vec![], vec![]);
        let assembler = Assembler::new(&mig);
        assert!(!assembler.config.skip_unknown_segments);
    }

    #[test]
    fn test_skip_does_not_consume_nested_group_entry() {
        // Skip must NOT consume segments that are nested group entries.
        // SG4 expects [IDE, STS], nested SG5 expects [LOC].
        // Input: IDE, FOO, STS, LOC. FOO should be skipped, LOC goes to SG5.
        let sg5 = make_mig_group("SG5", vec!["LOC"], vec![]);
        let sg4 = make_mig_group("SG4", vec!["IDE", "STS"], vec![sg5]);
        let mig = make_mig_schema(vec!["UNH"], vec![sg4]);

        let segments = vec![
            make_owned_seg("UNH", vec![vec!["001"]]),
            make_owned_seg("IDE", vec![vec!["24"]]),
            make_owned_seg("FOO", vec![vec!["unknown"]]),
            make_owned_seg("STS", vec![vec!["7"]]),
            make_owned_seg("LOC", vec![vec!["Z16"]]),
        ];

        let assembler = Assembler::with_config(
            &mig,
            AssemblerConfig {
                skip_unknown_segments: true,
                ..Default::default()
            },
        );
        let tree = assembler.assemble_generic(&segments).unwrap();
        let sg4 = &tree.groups[0];
        let inst = &sg4.repetitions[0];

        // IDE + STS consumed, FOO skipped
        assert_eq!(inst.segments.len(), 2);
        assert_eq!(inst.segments[0].tag, "IDE");
        assert_eq!(inst.segments[1].tag, "STS");
        assert_eq!(inst.skipped_segments.len(), 1);
        assert_eq!(inst.skipped_segments[0].tag, "FOO");

        // LOC went to nested SG5
        assert_eq!(inst.child_groups.len(), 1);
        assert_eq!(inst.child_groups[0].group_id, "SG5");
        assert_eq!(inst.child_groups[0].repetitions[0].segments[0].tag, "LOC");
    }

    #[test]
    fn test_roundtrip_with_skip() {
        // Full roundtrip: assemble with skip → disassemble → byte-identical
        // including skipped segments in the output.
        use crate::disassembler::Disassembler;
        use crate::renderer::render_edifact;

        let sg8 = make_mig_group("SG8", vec!["SEQ", "CCI"], vec![]);
        let mig = make_mig_schema(vec!["UNH", "UNT"], vec![sg8]);

        let segments = vec![
            make_owned_seg("UNH", vec![vec!["001"]]),
            make_owned_seg("SEQ", vec![vec!["Z98"]]),
            make_owned_seg("RFF", vec![vec!["Z38", "REF1"]]),
            make_owned_seg("CCI", vec![vec!["Z30"]]),
            make_owned_seg("UNT", vec![vec!["4", "001"]]),
        ];

        let assembler = Assembler::with_config(
            &mig,
            AssemblerConfig {
                skip_unknown_segments: true,
                ..Default::default()
            },
        );
        let tree = assembler.assemble_generic(&segments).unwrap();

        let disassembler = Disassembler::new(&mig);
        let dis = disassembler.disassemble(&tree);
        let delimiters = edifact_primitives::EdifactDelimiters::default();
        let rendered = render_edifact(&dis, &delimiters);

        // All 5 segments should appear in output (including skipped RFF).
        // Disassembler emits MIG-guided segments first (SEQ, CCI),
        // then skipped segments (RFF) — so order within the group differs
        // from the original input, but all content is preserved.
        assert_eq!(dis.len(), 5);
        assert_eq!(dis[0].tag, "UNH");
        assert_eq!(dis[1].tag, "SEQ");
        assert_eq!(dis[2].tag, "CCI");
        assert_eq!(dis[3].tag, "RFF"); // skipped → emitted after MIG segments
        assert_eq!(dis[4].tag, "UNT");

        // Rendered output contains all segments
        assert!(rendered.contains("UNH+001"));
        assert!(rendered.contains("SEQ+Z98"));
        assert!(rendered.contains("RFF+Z38:REF1"));
        assert!(rendered.contains("CCI+Z30"));
        assert!(rendered.contains("UNT+4:001"));
    }

    // ── Variant-aware assembly tests ──

    #[test]
    fn test_variant_groups_interleaved_reps() {
        // Two SG8 variant definitions: one for SEQ+ZD7, one for SEQ+Z98.
        // Input has interleaved reps: ZD7, Z98, ZD7, Z98.
        // All should be collected into one SG8 group with 4 reps.
        let sg8_zd7 = make_mig_group_with_variant("SG8", vec!["SEQ", "CCI"], vec![], "ZD7");
        let sg8_z98 = make_mig_group_with_variant("SG8", vec!["SEQ", "RFF"], vec![], "Z98");

        let mig = make_mig_schema(vec!["UNH"], vec![sg8_zd7, sg8_z98]);

        let segments = vec![
            make_owned_seg("UNH", vec![vec!["001"]]),
            make_owned_seg("SEQ", vec![vec!["ZD7"]]),
            make_owned_seg("CCI", vec![vec!["Z30"]]),
            make_owned_seg("SEQ", vec![vec!["Z98"]]),
            make_owned_seg("RFF", vec![vec!["Z38", "REF1"]]),
            make_owned_seg("SEQ", vec![vec!["ZD7"]]),
            make_owned_seg("CCI", vec![vec!["Z31"]]),
            make_owned_seg("SEQ", vec![vec!["Z98"]]),
            make_owned_seg("RFF", vec![vec!["Z38", "REF2"]]),
        ];

        let assembler = Assembler::new(&mig);
        let result = assembler.assemble_generic(&segments).unwrap();

        assert_eq!(result.segments.len(), 1); // UNH
        assert_eq!(result.groups.len(), 1); // One combined SG8
        let sg8 = &result.groups[0];
        assert_eq!(sg8.group_id, "SG8");
        assert_eq!(sg8.repetitions.len(), 4);

        // ZD7 reps have SEQ+CCI, Z98 reps have SEQ+RFF
        assert_eq!(sg8.repetitions[0].segments[0].elements[0][0], "ZD7");
        assert_eq!(sg8.repetitions[0].segments[1].tag, "CCI");
        assert_eq!(sg8.repetitions[1].segments[0].elements[0][0], "Z98");
        assert_eq!(sg8.repetitions[1].segments[1].tag, "RFF");
        assert_eq!(sg8.repetitions[2].segments[0].elements[0][0], "ZD7");
        assert_eq!(sg8.repetitions[3].segments[0].elements[0][0], "Z98");
    }

    #[test]
    fn test_variant_groups_single_variant_type() {
        // Only Z98 reps, no ZD7 — still works with variant matching
        let sg8_zd7 = make_mig_group_with_variant("SG8", vec!["SEQ", "CCI"], vec![], "ZD7");
        let sg8_z98 = make_mig_group_with_variant("SG8", vec!["SEQ", "RFF"], vec![], "Z98");

        let mig = make_mig_schema(vec!["UNH"], vec![sg8_zd7, sg8_z98]);

        let segments = vec![
            make_owned_seg("UNH", vec![vec!["001"]]),
            make_owned_seg("SEQ", vec![vec!["Z98"]]),
            make_owned_seg("RFF", vec![vec!["Z38", "REF1"]]),
            make_owned_seg("SEQ", vec![vec!["Z98"]]),
            make_owned_seg("RFF", vec![vec!["Z38", "REF2"]]),
        ];

        let assembler = Assembler::new(&mig);
        let result = assembler.assemble_generic(&segments).unwrap();

        assert_eq!(result.groups.len(), 1);
        assert_eq!(result.groups[0].repetitions.len(), 2);
        assert_eq!(
            result.groups[0].repetitions[0].segments[0].elements[0][0],
            "Z98"
        );
        assert_eq!(
            result.groups[0].repetitions[1].segments[0].elements[0][0],
            "Z98"
        );
    }

    #[test]
    fn test_non_variant_groups_unchanged() {
        // Groups without variant_code behave exactly as before
        let sg2 = make_mig_group("SG2", vec!["NAD"], vec![]);
        let sg4 = make_mig_group("SG4", vec!["IDE", "STS"], vec![]);

        let mig = make_mig_schema(vec!["UNH", "BGM"], vec![sg2, sg4]);

        let segments = vec![
            make_owned_seg("UNH", vec![vec!["001"]]),
            make_owned_seg("BGM", vec![vec!["E01"]]),
            make_owned_seg("NAD", vec![vec!["MS", "9900123"]]),
            make_owned_seg("NAD", vec![vec!["MR", "9900456"]]),
            make_owned_seg("IDE", vec![vec!["24", "TX001"]]),
            make_owned_seg("STS", vec![vec!["7"], vec!["Z33"]]),
        ];

        let assembler = Assembler::new(&mig);
        let result = assembler.assemble_generic(&segments).unwrap();

        assert_eq!(result.segments.len(), 2);
        assert_eq!(result.groups.len(), 2);
        assert_eq!(result.groups[0].group_id, "SG2");
        assert_eq!(result.groups[0].repetitions.len(), 2);
        assert_eq!(result.groups[1].group_id, "SG4");
        assert_eq!(result.groups[1].repetitions.len(), 1);
    }

    #[test]
    fn test_variant_groups_with_nested_children() {
        // Variant groups can have nested child groups
        let sg10 = make_mig_group("SG10", vec!["CCI", "CAV"], vec![]);
        let sg8_zd7 = make_mig_group_with_variant("SG8", vec!["SEQ"], vec![sg10.clone()], "ZD7");
        let sg8_z98 = make_mig_group_with_variant("SG8", vec!["SEQ"], vec![sg10], "Z98");

        let mig = make_mig_schema(vec!["UNH"], vec![sg8_zd7, sg8_z98]);

        let segments = vec![
            make_owned_seg("UNH", vec![vec!["001"]]),
            make_owned_seg("SEQ", vec![vec!["ZD7"]]),
            make_owned_seg("CCI", vec![vec!["Z30"]]),
            make_owned_seg("CAV", vec![vec!["Z91", "Y"]]),
            make_owned_seg("SEQ", vec![vec!["Z98"]]),
            make_owned_seg("CCI", vec![vec!["Z31"]]),
            make_owned_seg("CAV", vec![vec!["Z91", "N"]]),
        ];

        let assembler = Assembler::new(&mig);
        let result = assembler.assemble_generic(&segments).unwrap();

        assert_eq!(result.groups.len(), 1);
        let sg8 = &result.groups[0];
        assert_eq!(sg8.repetitions.len(), 2);

        // First rep (ZD7) has nested SG10
        assert_eq!(sg8.repetitions[0].child_groups.len(), 1);
        assert_eq!(sg8.repetitions[0].child_groups[0].group_id, "SG10");
        assert_eq!(
            sg8.repetitions[0].child_groups[0].repetitions[0].segments[0].elements[0][0],
            "Z30"
        );

        // Second rep (Z98) has nested SG10
        assert_eq!(sg8.repetitions[1].child_groups.len(), 1);
        assert_eq!(
            sg8.repetitions[1].child_groups[0].repetitions[0].segments[0].elements[0][0],
            "Z31"
        );
    }

    #[test]
    fn test_variant_qualifier_check_prevents_wrong_variant_consumption() {
        // try_consume_group with variant_code set should NOT consume a segment
        // whose qualifier doesn't match, even if the tag matches.
        let sg8_zd7 = make_mig_group_with_variant("SG8", vec!["SEQ", "CCI"], vec![], "ZD7");

        let mig = make_mig_schema(vec!["UNH"], vec![sg8_zd7]);

        let segments = vec![
            make_owned_seg("UNH", vec![vec!["001"]]),
            make_owned_seg("SEQ", vec![vec!["Z98"]]), // Wrong qualifier
            make_owned_seg("CCI", vec![vec!["Z30"]]),
        ];

        let assembler = Assembler::new(&mig);
        let result = assembler.assemble_generic(&segments).unwrap();

        // SG8 should have no reps because Z98 != ZD7
        assert!(result.groups.is_empty());
    }

    #[test]
    fn test_mixed_variant_and_non_variant_groups() {
        // SG2 (no variant), then variant SG8s, then SG12 (no variant)
        let sg2 = make_mig_group("SG2", vec!["NAD"], vec![]);
        let sg8_zd7 = make_mig_group_with_variant("SG8", vec!["SEQ", "CCI"], vec![], "ZD7");
        let sg8_z98 = make_mig_group_with_variant("SG8", vec!["SEQ", "RFF"], vec![], "Z98");
        let sg12 = make_mig_group("SG12", vec!["NAD"], vec![]);

        let mig = make_mig_schema(vec!["UNH"], vec![sg2, sg8_zd7, sg8_z98, sg12]);

        let segments = vec![
            make_owned_seg("UNH", vec![vec!["001"]]),
            make_owned_seg("NAD", vec![vec!["MS", "9900123"]]),
            make_owned_seg("SEQ", vec![vec!["ZD7"]]),
            make_owned_seg("CCI", vec![vec!["Z30"]]),
            make_owned_seg("SEQ", vec![vec!["Z98"]]),
            make_owned_seg("RFF", vec![vec!["Z38", "REF1"]]),
            make_owned_seg("NAD", vec![vec!["Z65", "ID001"]]),
        ];

        let assembler = Assembler::new(&mig);
        let result = assembler.assemble_generic(&segments).unwrap();

        assert_eq!(result.groups.len(), 3); // SG2, SG8 (combined), SG12
        assert_eq!(result.groups[0].group_id, "SG2");
        assert_eq!(result.groups[0].repetitions.len(), 1);
        assert_eq!(result.groups[1].group_id, "SG8");
        assert_eq!(result.groups[1].repetitions.len(), 2);
        assert_eq!(result.groups[2].group_id, "SG12");
        assert_eq!(result.groups[2].repetitions.len(), 1);
    }

    #[test]
    fn test_assembler_disambiguates_shared_qualifier_by_full_code_profile() {
        // PID 55035 PIA variants: several mig slots share the primary qualifier
        // 4347='5' but differ at C212/7143 (one allows Z12, another SRW). With
        // only per-mig qualifier_map, the assembler consumes at the first matching
        // slot regardless of the composite code — the downstream validator is
        // then forced to second-guess the variant choice. Disambiguate at
        // assembly time by checking all code-bearing positions declared on the
        // MIG segment.
        use mig_types::schema::common::CodeDefinition;
        use mig_types::schema::mig::{MigComposite, MigDataElement};
        use std::collections::HashMap;

        fn code(value: &str) -> CodeDefinition {
            CodeDefinition {
                value: value.to_string(),
                name: value.to_string(),
                description: None,
            }
        }

        fn pia_slot(number: &str, composite_code: &str) -> MigSegment {
            MigSegment {
                id: "PIA".to_string(),
                name: "PIA".to_string(),
                description: None,
                counter: None,
                level: 1,
                number: Some(number.to_string()),
                max_rep_std: 1,
                max_rep_spec: 1,
                status_std: Some("M".to_string()),
                status_spec: Some("M".to_string()),
                example: None,
                data_elements: vec![MigDataElement {
                    id: "4347".to_string(),
                    name: "Produkt-ID-Funktion".to_string(),
                    description: None,
                    status_std: Some("M".to_string()),
                    status_spec: Some("M".to_string()),
                    format_std: None,
                    format_spec: None,
                    codes: vec![code("5")],
                    position: 0,
                }],
                composites: vec![MigComposite {
                    id: "C212".to_string(),
                    name: "Item Identifier".to_string(),
                    description: None,
                    status_std: Some("M".to_string()),
                    status_spec: Some("M".to_string()),
                    data_elements: vec![MigDataElement {
                        id: "7143".to_string(),
                        name: "Artikel/Dienstleistung-ID".to_string(),
                        description: None,
                        status_std: Some("M".to_string()),
                        status_spec: Some("M".to_string()),
                        format_std: None,
                        format_spec: None,
                        codes: vec![code(composite_code)],
                        position: 0,
                    }],
                    position: 1,
                }],
            }
        }

        let sg4 = MigSegmentGroup {
            segments: vec![
                crate::test_support::make_mig_segment_numbered("IDE", "00020"),
                pia_slot("00108", "Z12"),
                pia_slot("00197", "SRW"),
            ],
            ..make_mig_group("SG4", vec![], vec![])
        };
        let mig = make_mig_schema(vec!["UNH"], vec![sg4]);

        let segments = vec![
            make_owned_seg("UNH", vec![vec!["001"]]),
            make_owned_seg("IDE", vec![vec!["24"]]),
            // PIA+5+:::SRW — composite element 1, component 0 = "SRW"
            make_owned_seg("PIA", vec![vec!["5"], vec!["SRW"]]),
        ];

        // Both PIA mig slots share the qualifier (0,0)='5'. Without full-profile
        // matching, the first slot (00108) wins and the SRW composite is
        // mis-assigned to the Z12 variant.
        let mut qualifier_map = HashMap::new();
        qualifier_map.insert("00108".to_string(), (0, 0, "5".to_string()));
        qualifier_map.insert("00197".to_string(), (0, 0, "5".to_string()));

        let config = AssemblerConfig {
            skip_unknown_segments: false,
            qualifier_map,
            strict_code_matching: true,
        };
        let assembler = Assembler::with_config(&mig, config);
        let tree = assembler.assemble_generic(&segments).unwrap();

        let sg4_instance = &tree.groups[0].repetitions[0];
        let pia = sg4_instance
            .segments
            .iter()
            .find(|s| s.tag == "PIA")
            .expect("PIA consumed into SG4");
        assert_eq!(
            pia.mig_number.as_deref(),
            Some("00197"),
            "PIA+5+:::SRW must be assigned the SRW variant (mig=00197), not the Z12 variant"
        );
    }
}