"""Brain tool-calling loop with a fake adapter — no network, no real LLM."""
import brain


class FakeAdapter(brain.LLMAdapter):
    """Replays scripted event lists, one per stream() call (per tool round)."""

    def __init__(self, rounds):
        self._rounds = list(rounds)
        self.seen_messages = []

    async def stream(self, messages, tools):
        # snapshot the messages the brain sent this round
        self.seen_messages.append([dict(m) for m in messages])
        for ev in self._rounds.pop(0):
            yield ev


def _text(s):
    return {"type": "text", "text": s}


async def test_plain_text_turn_streams_sentences_and_commits_history():
    rounds = [[
        _text("Right away, sir. "),
        _text("It is done."),
        {"type": "finish", "reason": "stop"},
    ]]
    b = brain.Brain(FakeAdapter(rounds), tool_specs=None, tool_exec=None)
    messages = [{"role": "system", "content": "sys"},
                {"role": "user", "content": "do it"}]
    out = [s async for s in b.run(messages)]
    assert out == ["Right away, sir.", "It is done."]
    # final assistant message appended to the shared list
    assert messages[-1]["role"] == "assistant"
    assert "It is done." in messages[-1]["content"]


async def test_tool_call_round_executes_then_answers():
    calls = []

    async def tool_exec(name, args):
        calls.append((name, args))
        return "Ahmed likes flat whites"

    rounds = [
        [  # round 1: the model asks for a tool
            {"type": "tool_calls", "calls": [
                {"id": "c1", "name": "memory_search",
                 "arguments": '{"query": "coffee"}'}]},
            {"type": "finish", "reason": "tool_calls"},
        ],
        [  # round 2: the model answers using the tool result
            _text("You like flat whites, sir."),
            {"type": "finish", "reason": "stop"},
        ],
    ]
    adapter = FakeAdapter(rounds)
    b = brain.Brain(adapter, tool_specs=[{"x": 1}], tool_exec=tool_exec)
    messages = [{"role": "system", "content": "sys"},
                {"role": "user", "content": "how do I like coffee"}]
    out = [s async for s in b.run(messages)]

    assert calls == [("memory_search", {"query": "coffee"})]
    assert out == ["You like flat whites, sir."]

    roles = [m["role"] for m in messages]
    assert "tool" in roles                      # tool result recorded
    # an assistant tool_calls message precedes the tool result
    tool_idx = roles.index("tool")
    assert messages[tool_idx - 1]["role"] == "assistant"
    assert messages[tool_idx - 1]["tool_calls"][0]["function"]["name"] \
        == "memory_search"
    assert messages[tool_idx]["content"] == "Ahmed likes flat whites"
    # round 2 saw the tool result in its message history
    assert any(m.get("role") == "tool"
               for m in adapter.seen_messages[1])


async def test_bad_tool_arguments_do_not_crash():
    async def tool_exec(name, args):
        assert args == {}          # malformed JSON -> empty dict
        return "ok"

    rounds = [
        [{"type": "tool_calls", "calls": [
            {"id": "c1", "name": "tasks_list", "arguments": "{not json"}]},
         {"type": "finish", "reason": "tool_calls"}],
        [_text("Done."), {"type": "finish", "reason": "stop"}],
    ]
    b = brain.Brain(FakeAdapter(rounds), tool_specs=[{}], tool_exec=tool_exec)
    messages = [{"role": "user", "content": "x"}]
    out = [s async for s in b.run(messages)]
    assert out == ["Done."]