"""Tests for the parallel-streaming curriculum building blocks (Phase 3)."""
from app.agents.brain_agent import (
    BrainAgent,
    _ExpansionChild,
    _RootAndSections,
    _SectionExpansion,
    _SectionItem,
)


class _CapturingClient:
    """Captures the messages passed for the multi-doc prompt assertion."""

    def __init__(self):
        self.messages = None

    def structured_complete(self, messages, output_model, model=None):
        self.messages = messages
        return _RootAndSections(root_label="X", sections=[_SectionItem(label="S", source_docs=[1])])


def test_multi_doc_prompt_lists_documents_and_section_tagging():
    cap = _CapturingClient()
    brain = BrainAgent(client=cap)
    rs = brain.derive_root_and_sections("structure", "high_school", "", "", ["paperA.pdf", "paperB.pdf"])
    blob = " ".join(m["content"] for m in cap.messages)
    assert "paperA.pdf" in blob and "paperB.pdf" in blob
    assert "source_docs" in blob
    assert rs.sections[0].source_docs == [1]


def test_single_doc_prompt_omits_doc_listing():
    cap = _CapturingClient()
    BrainAgent(client=cap).derive_root_and_sections("structure", "high_school", "", "", ["only.pdf"])
    blob = " ".join(m["content"] for m in cap.messages)
    assert "source_docs" not in blob  # no multi-doc tagging instruction for a single paper


class _FakeClient:
    """Returns canned structured output based on the requested model."""

    def structured_complete(self, messages, output_model, model=None):
        if output_model is _RootAndSections:
            return _RootAndSections(
                root_label="Optimization",
                root_description="Methods to minimize loss",
                sections=[_SectionItem(label="Gradient Descent"), _SectionItem(label="Adam")],
            )
        if output_model is _SectionExpansion:
            return _SectionExpansion(
                children=[_ExpansionChild(label="Learning Rate"), _ExpansionChild(label="Momentum")]
            )
        raise AssertionError(f"unexpected model {output_model}")


def test_derive_root_and_sections():
    brain = BrainAgent(client=_FakeClient())
    rs = brain.derive_root_and_sections("doc structure", "high_school", "Optimization")
    assert rs.root_label == "Optimization"
    assert [s.label for s in rs.sections] == ["Gradient Descent", "Adam"]


def test_expand_section():
    brain = BrainAgent(client=_FakeClient())
    doc_excerpts = [{"index": 0, "filename": "doc.pdf", "structure_text": "doc structure"}]
    exp = brain.expand_section("Gradient Descent", doc_excerpts, "high_school")
    assert len(exp.children) == 2
    assert exp.children[0].label == "Learning Rate"
    # relationship defaults are valid
    assert exp.children[0].relationship in {"prerequisite", "related", "builds-on"}