from memabra.candidate_types import CandidateObject
from memabra.retrieval import CandidateRetriever, InMemoryCandidateProvider
from memabra.router import TaskContext


def test_retriever_ranks_trigger_matches_first():
    retriever = CandidateRetriever(
        [
            InMemoryCandidateProvider(
                candidate_type="memory",
                candidates=[
                    CandidateObject(
                        id="mem-weak",
                        type="memory",
                        title="Generic preference",
                        summary="A weak preference record",
                        confidence=0.4,
                        success_rate=0.4,
                        freshness=0.4,
                        triggers=["generic"],
                    ),
                    CandidateObject(
                        id="mem-strong",
                        type="memory",
                        title="Formatting preference",
                        summary="Telegram prefers plain text",
                        confidence=0.8,
                        success_rate=0.9,
                        freshness=0.9,
                        triggers=["telegram", "formatting"],
                        tags=["output"],
                    ),
                ],
            )
        ]
    )

    result = retriever.retrieve(
        TaskContext(user_input="Use my telegram formatting preference for the output."),
        top_k=2,
    )

    assert [candidate.id for candidate in result.memory] == ["mem-strong", "mem-weak"]
    assert result.skill == []
    assert result.tool == []