import json

import pytest

from ollie_assistants.assistant.agent_loop import AgentLoop
from ollie_assistants.assistant.tools import default_tool_registry
from ollie_assistants.llm.types import ChatConfig, ChatMessage, ChatResult, Role
from ollie_assistants.memory.store import InMemoryStore
from ollie_assistants.memory.tools import MemoryTools
from ollie_assistants.safety.guardrails import GuardrailService


class CapturingProvider:
    provider_id = "test"

    def __init__(self) -> None:
        self.messages = []
        self.calls = []
        self.main_messages = []

    async def chat(self, messages, config: ChatConfig) -> ChatResult:
        self.calls.append(list(messages))
        self.messages = list(messages)
        system_text = messages[0].content if messages and messages[0].role == Role.SYSTEM else ""
        if "should_search" in system_text:
            return ChatResult(
                text=json.dumps(
                    {
                        "should_search": True,
                        "query": "answer style preference",
                        "reason": "user asks remembered preference",
                    }
                )
            )
        if "should_write" in system_text:
            return ChatResult(
                text=json.dumps(
                    {
                        "should_write": True,
                        "content": "I prefer concise answers",
                        "kind": "preference",
                        "reason": "stable preference",
                    }
                )
            )
        self.main_messages = list(messages)
        return ChatResult(text="ok")


@pytest.mark.asyncio
async def test_agent_loop_executes_memory_tool_before_llm() -> None:
    guardrails = GuardrailService()
    provider = CapturingProvider()
    memory_tools = MemoryTools(InMemoryStore())
    memory_tools.remember_user_fact("s1", "I prefer concise answers", "setup", "preference")
    loop = AgentLoop(
        provider,
        default_tool_registry(memory_tools),
        guardrails,
    )
    safety = guardrails.check_input("what is my answer style preference?")

    result = await loop.run(
        run_id="run-1",
        session_id="s1",
        model="test",
        max_new_tokens=10,
        messages=[ChatMessage(Role.USER, "what is my answer style preference?")],
        user_message="what is my answer style preference?",
        safety_decision=safety,
    )

    tool_names = [tool.name for tool in result.tool_results]
    assert "search_memory" in tool_names
    assert "remember_user_fact" in tool_names
    assert all(message.role != Role.TOOL for message in provider.main_messages)
    assert any("Relevant memory" in message.content for message in provider.main_messages)
    assert any(event.name == "memory_search_decision" for event in result.events)
    assert any(event.name == "memory_write_decision" for event in result.events)