mig-assembly 0.1.30

//! 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>>,
}

/// 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>,
    /// 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,
}

/// 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) {
            let assembled = owned_to_assembled(seg);
            cursor.advance();
            Ok(Some(assembled))
        } else {
            Ok(None) // Segment not present (optional)
        }
    }

    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 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(),
                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)
                    for slot in &mig_group.segments[slot_idx..slot_idx + run_len] {
                        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: greedily consume all matching segments
                    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;
                }
            }

            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
}

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

#[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()]],
                }],
                child_groups: vec![],
                skipped_segments: vec![],
            }],
        };

        let sg4_instance = AssembledGroupInstance {
            segments: vec![
                AssembledSegment {
                    tag: "IDE".to_string(),
                    elements: vec![vec!["24".to_string(), "TX001".to_string()]],
                },
                AssembledSegment {
                    tag: "STS".to_string(),
                    elements: vec![vec!["7".to_string()]],
                },
            ],
            child_groups: vec![sg5],
            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");
    }

    // ── 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,
            },
        );
        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,
            },
        );
        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,
            },
        );
        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,
            },
        );
        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);
    }
}