from phase_z2_composition import CompositionUnit
from phase_z2_pipeline import (
    GENERIC_FALLBACK_FRAME_ID,
    GENERIC_FALLBACK_FRAME_TEMPLATE_ID,
    MdxSection,
    _collapse_units_to_parent_sections,
)


def _unit(source_ids, template="child_frame", score=1.0):
    return CompositionUnit(
        source_section_ids=list(source_ids),
        merge_type="single",
        frame_template_id=template,
        frame_id="frame",
        frame_number=1,
        confidence=0.9,
        label="light_edit",
        phase_z_status="matched_zone",
        raw_content="child raw",
        title="child title",
        score=score,
    )


def test_task14_child_sections_collapse_to_parent_units():
    parent_sections = [
        MdxSection("05-1", 1, "S1 title", "S1 raw"),
        MdxSection("05-2", 2, "S2 title", "S2 raw"),
    ]
    aligned_sections = [
        MdxSection("05-1", 1, "S1 title", "S1 raw"),
        MdxSection("05-2-sub-1", 2, "child 1", "child 1 raw"),
        MdxSection("05-2-sub-2", 2, "child 2", "child 2 raw"),
    ]
    units = [
        _unit(["05-1"], template="frame_a", score=0.5),
        _unit(["05-2-sub-1", "05-2-sub-2"], template="frame_b", score=0.7),
    ]

    collapsed, status_sections, audit = _collapse_units_to_parent_sections(
        units,
        aligned_sections,
        parent_sections,
        fallback_template_id=GENERIC_FALLBACK_FRAME_TEMPLATE_ID,
        fallback_frame_id=GENERIC_FALLBACK_FRAME_ID,
    )

    assert audit["applied"] is True
    assert [s.section_id for s in status_sections] == ["05-1", "05-2"]
    assert [u.source_section_ids for u in collapsed] == [["05-1"], ["05-2"]]
    assert collapsed[1].merge_type == "parent_section_unit"
    assert collapsed[1].title == "S2 title"
    assert collapsed[1].raw_content == "S2 raw"
    assert (
        collapsed[1].rationale["task14_parent_section_unit"]["aligned_child_section_ids"]
        == ["05-2-sub-1", "05-2-sub-2"]
    )


def test_task14_adds_parent_fallback_when_group_uncovered():
    parent_sections = [MdxSection("04-2", 2, "S2 title", "S2 raw")]
    aligned_sections = [
        MdxSection("04-2-sub-1", 2, "child 1", "child 1 raw"),
        MdxSection("04-2-sub-2", 2, "child 2", "child 2 raw"),
    ]

    collapsed, status_sections, audit = _collapse_units_to_parent_sections(
        [],
        aligned_sections,
        parent_sections,
        fallback_template_id=GENERIC_FALLBACK_FRAME_TEMPLATE_ID,
        fallback_frame_id=GENERIC_FALLBACK_FRAME_ID,
    )

    assert [s.section_id for s in status_sections] == ["04-2"]
    assert [u.source_section_ids for u in collapsed] == [["04-2"]]
    assert collapsed[0].merge_type == "parent_no_drop_fallback"
    assert audit["added_fallbacks"][0]["parent_section_id"] == "04-2"