mig_assembly/

disassembler.rs

//! Tree disassembler — converts AssembledTree back to ordered segments.
//!
//! Walks the MIG schema tree in order. For each MIG node that has
//! corresponding data in the assembled tree, emits segments in MIG order.

use crate::assembler::{AssembledGroup, AssembledGroupInstance, AssembledSegment, AssembledTree};
use mig_types::schema::mig::{MigSchema, MigSegmentGroup};

/// Output segment from disassembly (owned data, ready for rendering).
#[derive(Debug, Clone)]
pub struct DisassembledSegment {
    pub tag: String,
    pub elements: Vec<Vec<String>>,
}

/// MIG-guided disassembler — walks the MIG tree to emit segments in correct order.
pub struct Disassembler<'a> {
    mig: &'a MigSchema,
}

impl<'a> Disassembler<'a> {
    pub fn new(mig: &'a MigSchema) -> Self {
        Self { mig }
    }

    /// Disassemble a tree into ordered segments following MIG sequence.
    ///
    /// Emits segments in correct EDIFACT order:
    /// 1. Pre-group top-level segments (e.g., UNB, UNH, BGM, DTM)
    /// 2. Groups (recursively, in MIG order)
    /// 3. Post-group top-level segments (e.g., UNT, UNZ)
    ///
    /// Uses MIG-guided ordering: walks the MIG schema tree and looks up
    /// matching data in the assembled tree. This handles both assembler output
    /// (already in MIG order) and reverse-mapped trees (may be in arbitrary order).
    pub fn disassemble(&self, tree: &AssembledTree) -> Vec<DisassembledSegment> {
        let mut output = Vec::new();

        // 1. Emit pre-group segments in MIG order
        let pre_group = &tree.segments[..tree.post_group_start];
        let mut consumed = vec![false; pre_group.len()];
        for mig_seg in &self.mig.segments {
            if let Some(idx) = pre_group
                .iter()
                .enumerate()
                .position(|(i, s)| !consumed[i] && s.tag == mig_seg.id)
            {
                output.push(assembled_to_disassembled(&pre_group[idx]));
                consumed[idx] = true;
            }
        }

        // 2. Emit groups in MIG order (lookup by group ID with consumption tracking).
        //    Between groups, emit any inter-group root segments (e.g., UNS in MSCONS).
        //    For variant groups (consecutive same-ID with variant_code), collect all
        //    variant definitions and match each rep to its correct variant.
        let mut consumed_groups = vec![false; tree.groups.len()];
        let mut mig_group_idx = 0;
        while mig_group_idx < self.mig.segment_groups.len() {
            let mig_group = &self.mig.segment_groups[mig_group_idx];

            // Check if this starts a variant set
            if mig_group.variant_code.is_some() {
                let variant_count = self.mig.segment_groups[mig_group_idx..]
                    .iter()
                    .take_while(|g| g.id == mig_group.id && g.variant_code.is_some())
                    .count();
                let variant_defs =
                    &self.mig.segment_groups[mig_group_idx..mig_group_idx + variant_count];

                if let Some(idx) = tree
                    .groups
                    .iter()
                    .enumerate()
                    .position(|(i, g)| !consumed_groups[i] && g.group_id == mig_group.id)
                {
                    if let Some(inter_segs) = tree.inter_group_segments.get(&idx) {
                        for seg in inter_segs {
                            output.push(assembled_to_disassembled(seg));
                        }
                    }
                    self.emit_variant_group(&tree.groups[idx], variant_defs, &mut output);
                    consumed_groups[idx] = true;
                }
                mig_group_idx += variant_count;
            } else {
                if let Some(idx) = tree
                    .groups
                    .iter()
                    .enumerate()
                    .position(|(i, g)| !consumed_groups[i] && g.group_id == mig_group.id)
                {
                    // Emit any inter-group segments that precede this group
                    if let Some(inter_segs) = tree.inter_group_segments.get(&idx) {
                        for seg in inter_segs {
                            output.push(assembled_to_disassembled(seg));
                        }
                    }
                    self.emit_group(&tree.groups[idx], mig_group, &mut output);
                    consumed_groups[idx] = true;
                }
                mig_group_idx += 1;
            }
        }

        // 2b. Emit trailing inter-group segments (after the last group).
        //     For ORDERS, UNS+S comes after SG29 (detail/summary boundary).
        let trailing_idx = tree.groups.len();
        if let Some(inter_segs) = tree.inter_group_segments.get(&trailing_idx) {
            for seg in inter_segs {
                output.push(assembled_to_disassembled(seg));
            }
        }

        // 3. Emit post-group segments (e.g., UNT, UNZ)
        for seg in &tree.segments[tree.post_group_start..] {
            output.push(assembled_to_disassembled(seg));
        }

        output
    }

    fn emit_group(
        &self,
        group: &AssembledGroup,
        mig_group: &MigSegmentGroup,
        output: &mut Vec<DisassembledSegment>,
    ) {
        for instance in &group.repetitions {
            self.emit_group_instance(instance, mig_group, output);
        }
    }

    /// Emit a group whose reps come from multiple variant definitions.
    ///
    /// For each rep, determine which variant it belongs to by checking the
    /// entry segment's qualifier against each variant's `variant_code`.
    /// Then emit the rep using that variant's segment/nested-group ordering.
    fn emit_variant_group(
        &self,
        group: &AssembledGroup,
        variant_defs: &[MigSegmentGroup],
        output: &mut Vec<DisassembledSegment>,
    ) {
        for instance in &group.repetitions {
            let entry_qual = instance
                .segments
                .first()
                .and_then(|s| s.elements.first())
                .and_then(|e| e.first())
                .map(|v| v.as_str())
                .unwrap_or("");

            let variant_def = variant_defs
                .iter()
                .find(|v| {
                    v.variant_code
                        .as_deref()
                        .map(|vc| vc.eq_ignore_ascii_case(entry_qual))
                        .unwrap_or(false)
                })
                .unwrap_or(&variant_defs[0]); // fallback to first variant

            self.emit_group_instance(instance, variant_def, output);
        }
    }

    fn emit_group_instance(
        &self,
        instance: &AssembledGroupInstance,
        mig_group: &MigSegmentGroup,
        output: &mut Vec<DisassembledSegment>,
    ) {
        // Emit segments in MIG order using tag-based lookup with consumption tracking.
        // This handles both assembler output (in MIG order) and reverse-mapped trees
        // (may be in arbitrary order).
        //
        // After MIG-guided emission, any remaining unconsumed segments are appended.
        // This handles cases where the assembler captured more segments than the MIG
        // defines (e.g., 6 RFFs when the merged MIG only has 4 slots).
        let mut consumed = vec![false; instance.segments.len()];
        for mig_seg in &mig_group.segments {
            if let Some(idx) = instance
                .segments
                .iter()
                .enumerate()
                .position(|(i, s)| !consumed[i] && s.tag == mig_seg.id)
            {
                output.push(assembled_to_disassembled(&instance.segments[idx]));
                consumed[idx] = true;
            }
        }
        // Emit any remaining segments not matched by MIG slots
        for (i, seg) in instance.segments.iter().enumerate() {
            if !consumed[i] {
                output.push(assembled_to_disassembled(seg));
            }
        }

        // Re-emit skipped segments (unknown segments preserved for roundtrip fidelity)
        for skipped in &instance.skipped_segments {
            output.push(assembled_to_disassembled(skipped));
        }

        // Child groups — lookup by group ID with consumption tracking.
        // Applies variant-aware logic recursively for nested variant groups.
        let mut consumed_child = vec![false; instance.child_groups.len()];
        let mut nested_idx = 0;
        while nested_idx < mig_group.nested_groups.len() {
            let nested_mig = &mig_group.nested_groups[nested_idx];

            if nested_mig.variant_code.is_some() {
                let variant_count = mig_group.nested_groups[nested_idx..]
                    .iter()
                    .take_while(|g| g.id == nested_mig.id && g.variant_code.is_some())
                    .count();
                let variant_defs = &mig_group.nested_groups[nested_idx..nested_idx + variant_count];

                if let Some(idx) = instance
                    .child_groups
                    .iter()
                    .enumerate()
                    .position(|(i, g)| !consumed_child[i] && g.group_id == nested_mig.id)
                {
                    self.emit_variant_group(&instance.child_groups[idx], variant_defs, output);
                    consumed_child[idx] = true;
                }
                nested_idx += variant_count;
            } else {
                if let Some(idx) = instance
                    .child_groups
                    .iter()
                    .enumerate()
                    .position(|(i, g)| !consumed_child[i] && g.group_id == nested_mig.id)
                {
                    self.emit_group(&instance.child_groups[idx], nested_mig, output);
                    consumed_child[idx] = true;
                }
                nested_idx += 1;
            }
        }
    }
}

fn assembled_to_disassembled(seg: &AssembledSegment) -> DisassembledSegment {
    DisassembledSegment {
        tag: seg.tag.clone(),
        elements: seg.elements.clone(),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::assembler::{
        AssembledGroup, AssembledGroupInstance, AssembledSegment, AssembledTree,
    };
    use crate::test_support::{make_mig_group, make_mig_segment};
    use mig_types::schema::mig::MigSchema;

    #[test]
    fn test_disassemble_top_level_only() {
        let mig = 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: vec![make_mig_segment("UNH"), make_mig_segment("BGM")],
            segment_groups: vec![],
        };

        let tree = AssembledTree {
            segments: vec![
                AssembledSegment {
                    tag: "UNH".to_string(),
                    elements: vec![
                        vec!["1".to_string()],
                        vec![
                            "UTILMD".to_string(),
                            "D".to_string(),
                            "11A".to_string(),
                            "UN".to_string(),
                            "S2.1".to_string(),
                        ],
                    ],
                    mig_number: None,
                },
                AssembledSegment {
                    tag: "BGM".to_string(),
                    elements: vec![
                        vec!["E01".to_string()],
                        vec!["MSG001".to_string()],
                        vec!["9".to_string()],
                    ],
                    mig_number: None,
                },
            ],
            groups: vec![],
            post_group_start: 2,
            inter_group_segments: std::collections::BTreeMap::new(),
        };

        let disassembler = Disassembler::new(&mig);
        let segments = disassembler.disassemble(&tree);

        assert_eq!(segments.len(), 2);
        assert_eq!(segments[0].tag, "UNH");
        assert_eq!(segments[1].tag, "BGM");
        assert_eq!(segments[0].elements[0], vec!["1"]);
    }

    #[test]
    fn test_disassemble_with_groups() {
        let mig = MigSchema {
            message_type: "UTILMD".to_string(),
            variant: None,
            version: "S2.1".to_string(),
            publication_date: "".to_string(),
            author: "".to_string(),
            format_version: "FV2504".to_string(),
            source_file: "test".to_string(),
            segments: vec![make_mig_segment("UNH"), make_mig_segment("BGM")],
            segment_groups: vec![make_mig_group("SG2", vec!["NAD", "LOC"], vec![])],
        };

        let tree = AssembledTree {
            segments: vec![
                AssembledSegment {
                    tag: "UNH".to_string(),
                    elements: vec![vec!["1".to_string()]],
                    mig_number: None,
                },
                AssembledSegment {
                    tag: "BGM".to_string(),
                    elements: vec![vec!["E01".to_string()]],
                    mig_number: None,
                },
            ],
            post_group_start: 2,
            groups: vec![AssembledGroup {
                group_id: "SG2".to_string(),
                repetitions: vec![
                    AssembledGroupInstance {
                        segments: vec![AssembledSegment {
                            tag: "NAD".to_string(),
                            elements: vec![vec!["MS".to_string()]],
                            mig_number: None,
                        }],
                        child_groups: vec![],
                        entry_mig_number: None,
                        variant_mig_numbers: vec![],
                        skipped_segments: vec![],
                        skipped_positions: Vec::new(),
                    },
                    AssembledGroupInstance {
                        segments: vec![AssembledSegment {
                            tag: "NAD".to_string(),
                            elements: vec![vec!["MR".to_string()]],
                            mig_number: None,
                        }],
                        child_groups: vec![],
                        entry_mig_number: None,
                        variant_mig_numbers: vec![],
                        skipped_segments: vec![],
                        skipped_positions: Vec::new(),
                    },
                ],
            }],
            inter_group_segments: std::collections::BTreeMap::new(),
        };

        let disassembler = Disassembler::new(&mig);
        let segments = disassembler.disassemble(&tree);

        assert_eq!(segments.len(), 4); // UNH, BGM, NAD(MS), NAD(MR)
        assert_eq!(segments[0].tag, "UNH");
        assert_eq!(segments[1].tag, "BGM");
        assert_eq!(segments[2].tag, "NAD");
        assert_eq!(segments[2].elements[0][0], "MS");
        assert_eq!(segments[3].tag, "NAD");
        assert_eq!(segments[3].elements[0][0], "MR");
    }

    #[test]
    fn test_disassemble_nested_groups() {
        let sg3 = make_mig_group("SG3", vec!["CTA", "COM"], vec![]);
        let mig = MigSchema {
            message_type: "UTILMD".to_string(),
            variant: None,
            version: "S2.1".to_string(),
            publication_date: "".to_string(),
            author: "".to_string(),
            format_version: "FV2504".to_string(),
            source_file: "test".to_string(),
            segments: vec![make_mig_segment("UNH")],
            segment_groups: vec![make_mig_group("SG2", vec!["NAD"], vec![sg3])],
        };

        let tree = AssembledTree {
            segments: vec![AssembledSegment {
                tag: "UNH".to_string(),
                elements: vec![vec!["1".to_string()]],
                mig_number: None,
            }],
            post_group_start: 1,
            groups: vec![AssembledGroup {
                group_id: "SG2".to_string(),
                repetitions: vec![AssembledGroupInstance {
                    segments: vec![AssembledSegment {
                        tag: "NAD".to_string(),
                        elements: vec![vec!["MS".to_string()]],
                        mig_number: None,
                    }],
                    child_groups: vec![AssembledGroup {
                        group_id: "SG3".to_string(),
                        repetitions: vec![AssembledGroupInstance {
                            segments: vec![
                                AssembledSegment {
                                    tag: "CTA".to_string(),
                                    elements: vec![vec!["IC".to_string()]],
                                    mig_number: None,
                                },
                                AssembledSegment {
                                    tag: "COM".to_string(),
                                    elements: vec![vec![
                                        "040@ex.com".to_string(),
                                        "EM".to_string(),
                                    ]],
                                    mig_number: None,
                                },
                            ],
                            child_groups: vec![],
                            entry_mig_number: None,
                            variant_mig_numbers: vec![],
                            skipped_segments: vec![],
                            skipped_positions: Vec::new(),
                        }],
                    }],
                    entry_mig_number: None,
                    variant_mig_numbers: vec![],
                    skipped_segments: vec![],
                    skipped_positions: Vec::new(),
                }],
            }],
            inter_group_segments: std::collections::BTreeMap::new(),
        };

        let disassembler = Disassembler::new(&mig);
        let segments = disassembler.disassemble(&tree);

        assert_eq!(segments.len(), 4); // UNH, NAD, CTA, COM
        assert_eq!(segments[0].tag, "UNH");
        assert_eq!(segments[1].tag, "NAD");
        assert_eq!(segments[2].tag, "CTA");
        assert_eq!(segments[3].tag, "COM");
    }

    #[test]
    fn test_disassemble_variant_groups_uses_per_variant_ordering() {
        use crate::test_support::make_mig_group_with_variant;

        // Two SG8 variant definitions with different segment ordering:
        // - ZD7 variant: SEQ, RFF, DTM
        // - Z98 variant: SEQ, CCI, CAV
        let sg8_zd7 = make_mig_group_with_variant("SG8", vec!["SEQ", "RFF", "DTM"], vec![], "ZD7");
        let sg8_z98 = make_mig_group_with_variant("SG8", vec!["SEQ", "CCI", "CAV"], vec![], "Z98");

        let mig = MigSchema {
            message_type: "UTILMD".to_string(),
            variant: None,
            version: "S2.1".to_string(),
            publication_date: "".to_string(),
            author: "".to_string(),
            format_version: "FV2504".to_string(),
            source_file: "test".to_string(),
            segments: vec![make_mig_segment("UNH")],
            segment_groups: vec![sg8_zd7, sg8_z98],
        };

        // Assembled tree has 3 reps: ZD7, Z98, ZD7 (interleaved)
        let tree = AssembledTree {
            segments: vec![AssembledSegment {
                tag: "UNH".to_string(),
                elements: vec![vec!["1".to_string()]],
                mig_number: None,
            }],
            post_group_start: 1,
            groups: vec![AssembledGroup {
                group_id: "SG8".to_string(),
                repetitions: vec![
                    // Rep 0: ZD7 variant — has SEQ, DTM, RFF (out of MIG order)
                    AssembledGroupInstance {
                        segments: vec![
                            AssembledSegment {
                                tag: "SEQ".to_string(),
                                elements: vec![vec!["ZD7".to_string()]],
                                mig_number: None,
                            },
                            AssembledSegment {
                                tag: "DTM".to_string(),
                                elements: vec![vec!["303".to_string()]],
                                mig_number: None,
                            },
                            AssembledSegment {
                                tag: "RFF".to_string(),
                                elements: vec![vec!["Z13".to_string()]],
                                mig_number: None,
                            },
                        ],
                        child_groups: vec![],
                        entry_mig_number: None,
                        variant_mig_numbers: vec![],
                        skipped_segments: vec![],
                        skipped_positions: Vec::new(),
                    },
                    // Rep 1: Z98 variant — has SEQ, CAV, CCI (out of MIG order)
                    AssembledGroupInstance {
                        segments: vec![
                            AssembledSegment {
                                tag: "SEQ".to_string(),
                                elements: vec![vec!["Z98".to_string()]],
                                mig_number: None,
                            },
                            AssembledSegment {
                                tag: "CAV".to_string(),
                                elements: vec![vec!["Z91".to_string()]],
                                mig_number: None,
                            },
                            AssembledSegment {
                                tag: "CCI".to_string(),
                                elements: vec![vec!["".to_string()]],
                                mig_number: None,
                            },
                        ],
                        child_groups: vec![],
                        entry_mig_number: None,
                        variant_mig_numbers: vec![],
                        skipped_segments: vec![],
                        skipped_positions: Vec::new(),
                    },
                    // Rep 2: another ZD7 variant
                    AssembledGroupInstance {
                        segments: vec![
                            AssembledSegment {
                                tag: "SEQ".to_string(),
                                elements: vec![vec!["ZD7".to_string()]],
                                mig_number: None,
                            },
                            AssembledSegment {
                                tag: "RFF".to_string(),
                                elements: vec![vec!["Z34".to_string()]],
                                mig_number: None,
                            },
                        ],
                        child_groups: vec![],
                        entry_mig_number: None,
                        variant_mig_numbers: vec![],
                        skipped_segments: vec![],
                        skipped_positions: Vec::new(),
                    },
                ],
            }],
            inter_group_segments: std::collections::BTreeMap::new(),
        };

        let disassembler = Disassembler::new(&mig);
        let segments = disassembler.disassemble(&tree);

        // UNH + 3 reps
        assert_eq!(segments[0].tag, "UNH");

        // Rep 0 (ZD7): MIG order is SEQ, RFF, DTM
        assert_eq!(segments[1].tag, "SEQ");
        assert_eq!(segments[1].elements[0][0], "ZD7");
        assert_eq!(segments[2].tag, "RFF"); // reordered from position 3 to 2
        assert_eq!(segments[3].tag, "DTM"); // reordered from position 2 to 3

        // Rep 1 (Z98): MIG order is SEQ, CCI, CAV
        assert_eq!(segments[4].tag, "SEQ");
        assert_eq!(segments[4].elements[0][0], "Z98");
        assert_eq!(segments[5].tag, "CCI"); // reordered from position 3 to 2
        assert_eq!(segments[6].tag, "CAV"); // reordered from position 2 to 3

        // Rep 2 (ZD7): SEQ, RFF (no DTM)
        assert_eq!(segments[7].tag, "SEQ");
        assert_eq!(segments[7].elements[0][0], "ZD7");
        assert_eq!(segments[8].tag, "RFF");

        assert_eq!(segments.len(), 9);
    }

    #[test]
    fn test_disassemble_empty_tree() {
        let mig = MigSchema {
            message_type: "UTILMD".to_string(),
            variant: None,
            version: "S2.1".to_string(),
            publication_date: "".to_string(),
            author: "".to_string(),
            format_version: "FV2504".to_string(),
            source_file: "test".to_string(),
            segments: vec![make_mig_segment("UNH")],
            segment_groups: vec![],
        };

        let tree = AssembledTree {
            segments: vec![],
            groups: vec![],
            post_group_start: 0,
            inter_group_segments: std::collections::BTreeMap::new(),
        };

        let disassembler = Disassembler::new(&mig);
        let segments = disassembler.disassemble(&tree);
        assert!(segments.is_empty());
    }
}