"""Tests for Agent memory integration."""
import pytest
import textwrap

from logos_adk.agent import Agent
from logos_adk.memory import LongTermMemory, MemoryEntry, SessionMemory, SharedMemory
from logos_adk.memory.backends import InMemoryBackend, SQLiteBackend
from logos_adk.providers import provider_registry
from logos_adk.providers.mock import MockProvider
from logos_adk.runtime import _default_runtime_holder
from logos_adk.types import Message


@pytest.fixture(autouse=True)
def _reset_runtime_and_registry():
    _default_runtime_holder.reset()
    provider_registry._entries.clear()
    yield
    _default_runtime_holder.reset()
    provider_registry._entries.clear()


@pytest.mark.asyncio
async def test_agent_run_includes_retrieved_memory():
    backend = InMemoryBackend()
    remembered = LongTermMemory(backend=backend)
    seed_agent = Agent("seed").memory(remembered)
    seed_provider = MockProvider(responses=["Saved"])
    seed_agent.model(seed_provider)
    await seed_agent.run("I prefer tea", session_id="s1")

    provider = MockProvider(responses=["You prefer tea"])
    agent = Agent("bot").model(provider).memory(LongTermMemory(backend=backend))
    await agent.run("What do I prefer?", session_id="s1")

    messages = provider.call_history[0]["messages"]
    memory_messages = [message for message in messages if message.role == "system"]
    assert any("I prefer tea" in message.content for message in memory_messages)


@pytest.mark.asyncio
async def test_agent_run_commits_long_term_memory():
    backend = InMemoryBackend()
    provider = MockProvider(responses=["Hello back"])
    agent = Agent("bot").model(provider).memory(LongTermMemory(backend=backend))

    await agent.run("hello there", session_id="s1")

    entries = await backend.search("hello there", session_id="s1")
    assert len(entries) == 1
    assert "Assistant: Hello back" in entries[0].content


@pytest.mark.asyncio
async def test_session_memory_trims_agent_state():
    provider = MockProvider(responses=["one", "two"])
    agent = Agent("bot").model(provider).memory(SessionMemory(max_messages=2))

    await agent.run("first", session_id="s1")
    await agent.run("second", session_id="s1")

    state = agent._get_state("s1")
    assert [message.content for message in state.messages] == ["second", "two"]


@pytest.mark.asyncio
async def test_session_memory_summarize_injects_summary_next_turn():
    provider = MockProvider(responses=["one", "two", "three"])
    agent = Agent("bot").model(provider).memory(
        SessionMemory(strategy="summarize", max_messages=2)
    )

    await agent.run("first", session_id="s1")
    await agent.run("second", session_id="s1")
    await agent.run("third", session_id="s1")

    third_call_messages = provider.call_history[2]["messages"]
    assert third_call_messages[0].role == "system"
    assert third_call_messages[0].content.startswith("Session summary:\n")
    assert "first" in third_call_messages[0].content
    assert [message.content for message in third_call_messages[1:]] == [
        "second",
        "two",
        "third",
    ]


@pytest.mark.asyncio
async def test_session_memory_token_limit_injects_summary_next_turn():
    provider = MockProvider(responses=["one", "two", "three"])
    agent = Agent("bot").model(provider).memory(
        SessionMemory(strategy="token_limit", max_tokens=45)
    )

    await agent.run("alpha " * 40, session_id="s1")
    await agent.run("recent question", session_id="s1")
    await agent.run("third", session_id="s1")

    third_call_messages = provider.call_history[2]["messages"]
    assert third_call_messages[0].role == "system"
    assert third_call_messages[0].content.startswith("Session summary:\n")
    assert "alpha" in third_call_messages[0].content.lower()
    assert [message.content for message in third_call_messages[1:]] == [
        "one",
        "recent question",
        "two",
        "third",
    ]


@pytest.mark.asyncio
async def test_session_memory_token_limit_keeps_recent_context_next_turn():
    provider = MockProvider(responses=["one", "two", "three"])
    agent = Agent("bot").model(provider).memory(
        SessionMemory(strategy="token_limit", max_tokens=5)
    )

    await agent.run("alpha", session_id="s1")
    await agent.run("beta gamma", session_id="s1")
    await agent.run("delta", session_id="s1")

    third_call_messages = provider.call_history[2]["messages"]
    assert [message.content for message in third_call_messages] == [
        "beta gamma",
        "two",
        "delta",
    ]


@pytest.mark.asyncio
async def test_memory_is_injected_before_conversation():
    backend = InMemoryBackend()
    memory = LongTermMemory(backend=backend)
    await backend.store(
        MemoryEntry(
            content="The customer account is premium",
            session_id="s1",
        )
    )

    provider = MockProvider(responses=["Premium confirmed"])
    agent = Agent("bot").model(provider).memory(memory).system("Be concise")
    await agent.run("What account tier is this?", session_id="s1")

    messages = provider.call_history[0]["messages"]
    assert messages[0].content == "Be concise"
    assert messages[1].role == "system"
    assert "premium" in messages[1].content.lower()
    assert messages[-1] == Message(role="user", content="What account tier is this?")


@pytest.mark.asyncio
async def test_sqlite_long_term_memory_persists_across_agent_instances(tmp_path):
    db_path = tmp_path / "memory.db"

    seed_provider = MockProvider(responses=["Saved"])
    seed_agent = Agent("seed").model(seed_provider).memory(
        LongTermMemory(backend=SQLiteBackend(str(db_path)))
    )
    await seed_agent.run("My favorite drink is tea", session_id="s1")

    provider = MockProvider(responses=["You like tea"])
    agent = Agent("bot").model(provider).memory(
        LongTermMemory(backend=SQLiteBackend(str(db_path)))
    )
    await agent.run("What is my favorite drink?", session_id="s1")

    messages = provider.call_history[0]["messages"]
    memory_messages = [message for message in messages if message.role == "system"]
    assert any("favorite drink is tea" in message.content for message in memory_messages)


@pytest.mark.asyncio
async def test_sqlite_long_term_memory_respects_session_scope(tmp_path):
    db_path = tmp_path / "memory.db"
    backend = SQLiteBackend(str(db_path))

    provider = MockProvider(responses=["Saved A", "Saved B", "Unknown"])
    agent = Agent("bot").model(provider).memory(LongTermMemory(backend=backend))

    await agent.run("I am session A", session_id="a")
    await agent.run("I am session B", session_id="b")

    other_provider = MockProvider(responses=["Checking"])
    other_agent = Agent("bot2").model(other_provider).memory(
        LongTermMemory(backend=SQLiteBackend(str(db_path)))
    )
    await other_agent.run("session", session_id="a")

    messages = other_provider.call_history[0]["messages"]
    memory_messages = [message for message in messages if message.role == "system"]
    assert any("session A" in message.content for message in memory_messages)
    assert all("session B" not in message.content for message in memory_messages)


@pytest.mark.asyncio
async def test_shared_memory_shares_context_between_agents():
    backend = InMemoryBackend()
    shared = SharedMemory(backend=backend)

    writer_provider = MockProvider(responses=["Saved for everyone"])
    writer = Agent("writer").model(writer_provider).memory(shared)
    await writer.run("Project codename is Atlas", session_id="writer-session")

    reader_provider = MockProvider(responses=["Atlas confirmed"])
    reader = Agent("reader").model(reader_provider).memory(
        SharedMemory(backend=backend)
    )
    await reader.run("Atlas", session_id="reader-session")

    messages = reader_provider.call_history[0]["messages"]
    memory_messages = [message for message in messages if message.role == "system"]
    assert any("Project codename is Atlas" in message.content for message in memory_messages)


@pytest.mark.asyncio
async def test_shared_memory_ignores_session_scope_with_sqlite_backend(tmp_path):
    db_path = tmp_path / "shared-memory.db"

    researcher_provider = MockProvider(responses=["Stored"])
    researcher = Agent("researcher").model(researcher_provider).memory(
        SharedMemory(backend=SQLiteBackend(str(db_path)))
    )
    await researcher.run("Launch city is Berlin", session_id="alpha")

    planner_provider = MockProvider(responses=["Berlin noted"])
    planner = Agent("planner").model(planner_provider).memory(
        SharedMemory(backend=SQLiteBackend(str(db_path)))
    )
    await planner.run("Berlin", session_id="beta")

    messages = planner_provider.call_history[0]["messages"]
    memory_messages = [message for message in messages if message.role == "system"]
    assert any("Launch city is Berlin" in message.content for message in memory_messages)


@pytest.mark.asyncio
async def test_shared_memory_from_config_shares_between_agents(tmp_path):
    db_path = tmp_path / "shared.db"
    writer_config = tmp_path / "writer.yaml"
    reader_config = tmp_path / "reader.yaml"

    shared_yaml = textwrap.dedent(
        f"""
        memory:
          shared:
            backend: sqlite
            path: {db_path}
            retrieval: keyword
        """
    ).strip()

    writer_config.write_text(
        "\n".join(
            [
                "name: writer",
                "model: mock-writer",
                shared_yaml,
            ]
        )
    )
    reader_config.write_text(
        "\n".join(
            [
                "name: reader",
                "model: mock-reader",
                shared_yaml,
            ]
        )
    )

    writer_provider = MockProvider(responses=["Stored"])
    reader_provider = MockProvider(responses=["Retrieved"])

    writer = Agent.from_config(str(writer_config)).model(writer_provider)
    reader = Agent.from_config(str(reader_config)).model(reader_provider)

    await writer.run("Secret launch word is Aurora", session_id="writer-session")
    await reader.run("Aurora", session_id="reader-session")

    messages = reader_provider.call_history[0]["messages"]
    memory_messages = [message for message in messages if message.role == "system"]
    assert any("Aurora" in message.content for message in memory_messages)
    assert any("launch word" in message.content for message in memory_messages)
