//! Convert an [`AssembledTree`] to a flat list of [`OwnedSegment`]s.
//!
//! This is needed to build an [`EvaluationContext`] for AHB condition evaluation,
//! which requires flat `&[OwnedSegment]` input.

use mig_assembly::assembler::{AssembledGroup, AssembledTree};
use mig_types::segment::OwnedSegment;

/// Flatten an [`AssembledTree`] into a sequential list of [`OwnedSegment`]s.
///
/// Segments are emitted in EDIFACT order:
/// 1. Pre-group root segments (UNH, BGM, etc.)
/// 2. Groups (recursively flattened)
/// 3. Inter-group segments (UNS, etc.)
/// 4. Post-group root segments (UNT, etc.)
pub fn tree_to_owned_segments(tree: &AssembledTree) -> Vec<OwnedSegment> {
    let mut segments = Vec::new();
    let mut number = 1u32;

    // Pre-group root segments
    for seg in &tree.segments[..tree.post_group_start] {
        segments.push(OwnedSegment {
            id: seg.tag.clone(),
            elements: seg.elements.clone(),
            segment_number: number,
        });
        number += 1;
    }

    // Groups with inter-group segments
    for (i, group) in tree.groups.iter().enumerate() {
        // Emit inter-group segments that precede this group
        if let Some(inter_segs) = tree.inter_group_segments.get(&i) {
            for seg in inter_segs {
                segments.push(OwnedSegment {
                    id: seg.tag.clone(),
                    elements: seg.elements.clone(),
                    segment_number: number,
                });
                number += 1;
            }
        }
        collect_group_segments(group, &mut segments, &mut number);
    }

    // Post-group root segments (UNT, etc.)
    for seg in &tree.segments[tree.post_group_start..] {
        segments.push(OwnedSegment {
            id: seg.tag.clone(),
            elements: seg.elements.clone(),
            segment_number: number,
        });
        number += 1;
    }

    segments
}

/// Recursively collect segments from an assembled group and its repetitions.
fn collect_group_segments(
    group: &AssembledGroup,
    segments: &mut Vec<OwnedSegment>,
    number: &mut u32,
) {
    for rep in &group.repetitions {
        // Segments in this repetition
        for seg in &rep.segments {
            segments.push(OwnedSegment {
                id: seg.tag.clone(),
                elements: seg.elements.clone(),
                segment_number: *number,
            });
            *number += 1;
        }

        // Nested child groups
        for child_group in &rep.child_groups {
            collect_group_segments(child_group, segments, number);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use mig_assembly::assembler::{AssembledGroupInstance, AssembledSegment};
    use std::collections::BTreeMap;

    fn seg(tag: &str, elements: Vec<Vec<&str>>) -> AssembledSegment {
        AssembledSegment {
            tag: tag.to_string(),
            elements: elements
                .into_iter()
                .map(|e| e.into_iter().map(|s| s.to_string()).collect())
                .collect(),
        }
    }

    #[test]
    fn test_empty_tree() {
        let tree = AssembledTree {
            segments: vec![],
            groups: vec![],
            post_group_start: 0,
            inter_group_segments: BTreeMap::new(),
        };
        let result = tree_to_owned_segments(&tree);
        assert!(result.is_empty());
    }

    #[test]
    fn test_flat_segments_only() {
        let tree = AssembledTree {
            segments: vec![
                seg("UNH", vec![vec!["1"], vec!["UTILMD"]]),
                seg("UNT", vec![vec!["5"], vec!["1"]]),
            ],
            groups: vec![],
            post_group_start: 1,
            inter_group_segments: BTreeMap::new(),
        };
        let result = tree_to_owned_segments(&tree);
        assert_eq!(result.len(), 2);
        assert_eq!(result[0].id, "UNH");
        assert_eq!(result[0].segment_number, 1);
        assert_eq!(result[1].id, "UNT");
        assert_eq!(result[1].segment_number, 2);
    }

    #[test]
    fn test_with_groups() {
        let tree = AssembledTree {
            segments: vec![seg("UNH", vec![vec!["1"]]), seg("UNT", vec![vec!["5"]])],
            groups: vec![AssembledGroup {
                group_id: "SG2".to_string(),
                repetitions: vec![AssembledGroupInstance {
                    segments: vec![seg("NAD", vec![vec!["MS"]]), seg("CTA", vec![vec!["IC"]])],
                    child_groups: vec![],
                    skipped_segments: vec![],
                }],
            }],
            post_group_start: 1,
            inter_group_segments: BTreeMap::new(),
        };
        let result = tree_to_owned_segments(&tree);
        assert_eq!(result.len(), 4);
        assert_eq!(result[0].id, "UNH");
        assert_eq!(result[0].segment_number, 1);
        assert_eq!(result[1].id, "NAD");
        assert_eq!(result[1].segment_number, 2);
        assert_eq!(result[2].id, "CTA");
        assert_eq!(result[2].segment_number, 3);
        assert_eq!(result[3].id, "UNT");
        assert_eq!(result[3].segment_number, 4);
    }

    #[test]
    fn test_nested_groups() {
        let tree = AssembledTree {
            segments: vec![seg("UNH", vec![vec!["1"]])],
            groups: vec![AssembledGroup {
                group_id: "SG4".to_string(),
                repetitions: vec![AssembledGroupInstance {
                    segments: vec![seg("IDE", vec![vec!["24"]])],
                    child_groups: vec![AssembledGroup {
                        group_id: "SG5".to_string(),
                        repetitions: vec![AssembledGroupInstance {
                            segments: vec![seg("LOC", vec![vec!["Z16"]])],
                            child_groups: vec![],
                            skipped_segments: vec![],
                        }],
                    }],
                    skipped_segments: vec![],
                }],
            }],
            post_group_start: 1,
            inter_group_segments: BTreeMap::new(),
        };
        let result = tree_to_owned_segments(&tree);
        assert_eq!(result.len(), 3);
        assert_eq!(result[0].id, "UNH");
        assert_eq!(result[1].id, "IDE");
        assert_eq!(result[2].id, "LOC");
        // Verify sequential numbering
        for (i, seg) in result.iter().enumerate() {
            assert_eq!(seg.segment_number, (i + 1) as u32);
        }
    }
}
