from memabra.candidate_types import CandidateObject
from memabra.retrieval import CandidateRetriever, InMemoryCandidateProvider
from memabra.router import RuleBasedRouter, TaskContext
from memabra.runner import MemabraRunner
from memabra.schemas import SchemaRegistry


def test_runner_produces_valid_draft_trajectory():
    retriever = CandidateRetriever(
        [
            InMemoryCandidateProvider(
                candidate_type="memory",
                candidates=[
                    CandidateObject(
                        id="mem-1",
                        type="memory",
                        title="Output preference",
                        summary="Prefer plain text on Telegram.",
                        triggers=["telegram", "preference"],
                        confidence=0.9,
                        success_rate=0.8,
                        freshness=0.9,
                        tags=["output"],
                    )
                ],
            )
        ]
    )
    runner = MemabraRunner(retriever=retriever, router=RuleBasedRouter())

    trajectory = runner.run(
        context=TaskContext(
            user_input="Use my telegram preference for this answer.",
            conversation_summary="User often cares about output formatting.",
        ),
        channel="telegram",
        user_id="oza",
    )

    SchemaRegistry().validate_trajectory(trajectory)
    assert trajectory["decisions"][0]["decision_type"] == "inject_memory"
    assert trajectory["candidate_sets"]["memory"][0]["id"] == "mem-1"
    assert len(trajectory["events"]) == 3


def test_runner_injects_episodic_candidate_when_case_index_matches(tmp_path):
    from memabra.case_index import CaseIndex
    from memabra.persistence import PersistenceStore

    store = PersistenceStore(base_dir=tmp_path / "artifacts")
    retriever = CandidateRetriever(
        [
            InMemoryCandidateProvider(
                candidate_type="memory",
                candidates=[],
            ),
            InMemoryCandidateProvider(
                candidate_type="skill",
                candidates=[],
            ),
            InMemoryCandidateProvider(
                candidate_type="tool",
                candidates=[],
            ),
        ]
    )
    runner = MemabraRunner(retriever=retriever, router=RuleBasedRouter(), persistence_store=store)

    # First run creates a trajectory
    trajectory1 = runner.run(
        context=TaskContext(user_input="Hello world"),
        channel="local",
        persist=True,
    )

    # Build case index from the trajectory
    case_index = CaseIndex()
    case_index.add(trajectory1)

    # Second run with case index should inject an episodic candidate
    runner_with_case = MemabraRunner(
        retriever=retriever,
        router=RuleBasedRouter(),
        persistence_store=store,
        case_index=case_index,
    )
    trajectory2 = runner_with_case.run(
        context=TaskContext(user_input="Hello world"),
        channel="local",
        persist=True,
    )

    memory_candidates = trajectory2["candidate_sets"]["memory"]
    assert any(c["id"].startswith("episodic-") for c in memory_candidates)
    # With a persistence store, the runner should generate a rich episodic summary
    assert any("Task:" in c["summary"] for c in memory_candidates)