"""Tests for memory primitives."""
import pytest

from logos_adk.agent import Agent
from logos_adk.errors import ConfigError
from logos_adk.memory import LongTermMemory, MemoryEntry, SessionMemory, SharedMemory
from logos_adk.memory.backends import InMemoryBackend, ShardedMemoryBackend, SQLiteBackend
from logos_adk.state import AgentState
from logos_adk.types import Message, Response, Usage


async def _remember_top_content(
    backend,
    *,
    retrieval: str,
    query: str,
    entries: list[tuple[str, str | None]],
) -> str | None:
    for content, session_id in entries:
        await backend.store(MemoryEntry(content=content, session_id=session_id))

    results = await LongTermMemory(
        backend=backend,
        retrieval=retrieval,
        limit=1,
    ).remember(
        [Message(role="user", content=query)],
        AgentState(session_id="s1"),
        session_id="s1",
    )
    return results[0].content if results else None


def test_agent_builder_memory():
    agent = Agent("bot").memory(SessionMemory(), SessionMemory(max_messages=5))
    assert len(agent._config.memories) == 2


def test_session_memory_rejects_unknown_strategy():
    with pytest.raises(ConfigError, match="Unsupported session memory strategy"):
        SessionMemory(strategy="unknown")


def test_session_memory_token_limit_requires_max_tokens():
    with pytest.raises(ConfigError, match="requires max_tokens"):
        SessionMemory(strategy="token_limit")


def test_session_memory_rejects_invalid_max_messages_type():
    with pytest.raises(ConfigError, match="max_messages must be an integer"):
        SessionMemory(max_messages="many")  # type: ignore[arg-type]


def test_memory_entry_defaults():
    entry = MemoryEntry(content="remember this")
    assert entry.content == "remember this"
    assert entry.embedding is None
    assert entry.metadata == {}
    assert entry.session_id is None


async def test_in_memory_backend_roundtrip():
    backend = InMemoryBackend()
    first = MemoryEntry(content="alpha topic", session_id="s1")
    second = MemoryEntry(content="beta topic", session_id="s2")

    await backend.store(first)
    await backend.store(second)

    results = await backend.search("alpha", session_id="s1")
    assert [entry.id for entry in results] == [first.id]

    await backend.delete(first.id)
    assert await backend.search("alpha", session_id="s1") == []


async def test_sqlite_backend_roundtrip(tmp_path):
    backend = SQLiteBackend(str(tmp_path / "memory.db"))
    first = MemoryEntry(
        content="alpha topic",
        session_id="s1",
        metadata={"kind": "fact"},
    )
    second = MemoryEntry(content="beta topic", session_id="s2")

    await backend.store(first)
    await backend.store(second)

    results = await backend.search("alpha", session_id="s1")
    assert len(results) == 1
    assert results[0].id == first.id
    assert results[0].metadata["kind"] == "fact"

    await backend.delete(first.id)
    assert await backend.search("alpha", session_id="s1") == []


async def test_sharded_memory_backend_routes_by_session_id(tmp_path):
    backend = ShardedMemoryBackend(
        [
            SQLiteBackend(str(tmp_path / "shard-a.db")),
            SQLiteBackend(str(tmp_path / "shard-b.db")),
        ]
    )
    first = MemoryEntry(content="alpha topic", session_id="tenant-a")
    second = MemoryEntry(content="beta topic", session_id="tenant-b")

    await backend.store(first)
    await backend.store(second)

    alpha_results = await backend.search("alpha", session_id="tenant-a")
    beta_results = await backend.search("beta", session_id="tenant-b")

    assert [entry.content for entry in alpha_results] == ["alpha topic"]
    assert [entry.content for entry in beta_results] == ["beta topic"]


async def test_session_memory_trims_history():
    memory = SessionMemory(max_messages=2)
    state = AgentState(
        messages=[
            Message(role="user", content="first"),
            Message(role="assistant", content="one"),
            Message(role="user", content="second"),
            Message(role="assistant", content="two"),
        ]
    )

    await memory.commit(
        state.messages,
        Response(content="two", tool_calls=[], usage=Usage(1, 1), raw={}),
        state,
    )

    assert [message.content for message in state.messages] == ["second", "two"]


async def test_session_memory_summarizes_older_history():
    memory = SessionMemory(strategy="summarize", max_messages=2)
    state = AgentState(
        messages=[
            Message(role="user", content="My name is Alex"),
            Message(role="assistant", content="Nice to meet you"),
            Message(role="user", content="I live in Berlin"),
            Message(role="assistant", content="Got it"),
        ]
    )

    await memory.commit(
        state.messages,
        Response(content="Got it", tool_calls=[], usage=Usage(1, 1), raw={}),
        state,
    )

    assert len(state.messages) == 3
    assert state.messages[0].role == "system"
    assert state.messages[0].content.startswith("Session summary:\n")
    assert "My name is Alex" in state.messages[0].content
    assert [message.content for message in state.messages[1:]] == [
        "I live in Berlin",
        "Got it",
    ]


async def test_session_memory_rolls_forward_existing_summary():
    memory = SessionMemory(strategy="summarize", max_messages=2)
    state = AgentState(
        messages=[
            Message(
                role="system",
                content="Session summary:\n- User: My name is Alex\n- Assistant: Nice to meet you",
            ),
            Message(role="user", content="I live in Berlin"),
            Message(role="assistant", content="Got it"),
            Message(role="user", content="I work as an engineer"),
            Message(role="assistant", content="Understood"),
        ]
    )

    await memory.commit(
        state.messages,
        Response(content="Understood", tool_calls=[], usage=Usage(1, 1), raw={}),
        state,
    )

    assert len(state.messages) == 3
    assert state.messages[0].role == "system"
    assert state.messages[0].content.count("Session summary:") == 1
    assert "My name is Alex" in state.messages[0].content
    assert "I live in Berlin" in state.messages[0].content
    assert [message.content for message in state.messages[1:]] == [
        "I work as an engineer",
        "Understood",
    ]


async def test_session_memory_token_limit_builds_summary_for_old_context():
    memory = SessionMemory(strategy="token_limit", max_tokens=45)
    long_user = "alpha " * 40
    long_assistant = "beta " * 40
    state = AgentState(
        messages=[
            Message(role="user", content=long_user.strip()),
            Message(role="assistant", content=long_assistant.strip()),
            Message(role="user", content="recent question"),
            Message(role="assistant", content="recent answer"),
        ]
    )

    await memory.commit(
        state.messages,
        Response(content="recent answer", tool_calls=[], usage=Usage(1, 1), raw={}),
        state,
    )

    assert state.messages[0].role == "system"
    assert state.messages[0].content.startswith("Session summary:\n")
    assert "alpha" in state.messages[0].content.lower()
    assert "beta" in state.messages[0].content.lower()
    assert [message.content for message in state.messages[1:]] == [
        "recent question",
        "recent answer",
    ]


async def test_session_memory_token_limit_keeps_recent_tail():
    memory = SessionMemory(strategy="token_limit", max_tokens=5)
    state = AgentState(
        messages=[
            Message(role="user", content="alpha beta gamma"),
            Message(role="assistant", content="one"),
            Message(role="user", content="delta epsilon"),
            Message(role="assistant", content="two"),
        ]
    )

    await memory.commit(
        state.messages,
        Response(content="two", tool_calls=[], usage=Usage(1, 1), raw={}),
        state,
    )

    assert [message.content for message in state.messages] == [
        "delta epsilon",
        "two",
    ]


async def test_session_memory_token_limit_keeps_latest_message_even_if_oversized():
    memory = SessionMemory(strategy="token_limit", max_tokens=1)
    state = AgentState(
        messages=[
            Message(role="user", content="short"),
            Message(role="assistant", content="this response is intentionally long"),
        ]
    )

    await memory.commit(
        state.messages,
        Response(
            content="this response is intentionally long",
            tool_calls=[],
            usage=Usage(1, 1),
            raw={},
        ),
        state,
    )

    assert [message.content for message in state.messages] == [
        "this response is intentionally long"
    ]


async def test_long_term_memory_commit_and_remember():
    backend = InMemoryBackend()
    memory = LongTermMemory(backend=backend)
    state = AgentState(session_id="s1")
    messages = [Message(role="user", content="I like tea")]
    response = Response(content="Noted", tool_calls=[], usage=Usage(1, 1), raw={})

    await memory.commit(messages, response, state, session_id="s1")
    results = await memory.remember(messages, state, session_id="s1")

    assert len(results) == 1
    assert "I like tea" in results[0].content


async def test_shared_memory_ignores_session_id():
    backend = InMemoryBackend()
    memory = SharedMemory(backend=backend)
    state = AgentState()
    messages = [Message(role="user", content="project codename is Atlas")]
    response = Response(content="Saved", tool_calls=[], usage=Usage(1, 1), raw={})

    await memory.commit(messages, response, state, session_id="alpha")
    results = await memory.remember(
        [Message(role="user", content="Atlas")],
        state,
        session_id="beta",
    )

    assert len(results) == 1
    assert results[0].session_id is None


def test_long_term_memory_rejects_unknown_retrieval():
    with pytest.raises(ConfigError, match="Unsupported long-term memory retrieval"):
        LongTermMemory(backend=InMemoryBackend(), retrieval="vector")


async def test_long_term_memory_hybrid_retrieval_ranks_token_overlap():
    backend = InMemoryBackend()
    memory = LongTermMemory(backend=backend, retrieval="hybrid", limit=2)
    state = AgentState(session_id="s1")

    await backend.store(
        MemoryEntry(
            content="The launch city is Berlin and the timeline is final",
            session_id="s1",
        )
    )
    await backend.store(
        MemoryEntry(
            content="The office has Berlin artwork",
            session_id="s1",
        )
    )

    results = await memory.remember(
        [Message(role="user", content="launch Berlin timeline")],
        state,
        session_id="s1",
    )

    assert len(results) >= 1
    assert "launch city is Berlin" in results[0].content


async def test_long_term_memory_semantic_retrieval_fuzzy_matches_query():
    backend = InMemoryBackend()
    memory = LongTermMemory(backend=backend, retrieval="semantic", limit=1)
    state = AgentState(session_id="s1")

    await backend.store(
        MemoryEntry(
            content="Customer prefers sparkling water over coffee",
            session_id="s1",
        )
    )

    results = await memory.remember(
        [Message(role="user", content="sparkly water preference")],
        state,
        session_id="s1",
    )

    assert len(results) == 1
    assert "sparkling water" in results[0].content


async def test_long_term_memory_hybrid_and_keyword_can_rank_differently():
    backend = InMemoryBackend()
    state = AgentState(session_id="s1")
    query = "launch plan berlin timeline"

    await backend.store(
        MemoryEntry(
            content="launch plan berlin",
            session_id="s1",
        )
    )
    await backend.store(
        MemoryEntry(
            content="berlin launch plan timeline final",
            session_id="s1",
        )
    )

    keyword_results = await LongTermMemory(
        backend=backend,
        retrieval="keyword",
        limit=1,
    ).remember([Message(role="user", content=query)], state, session_id="s1")
    hybrid_results = await LongTermMemory(
        backend=backend,
        retrieval="hybrid",
        limit=1,
    ).remember([Message(role="user", content=query)], state, session_id="s1")

    assert len(keyword_results) == 1
    assert len(hybrid_results) == 1
    assert keyword_results[0].content != ""
    assert hybrid_results[0].content == "berlin launch plan timeline final"


@pytest.mark.parametrize(
    ("retrieval", "query", "entries", "expected"),
    [
        (
            "keyword",
            "project codename atlas",
            [
                ("project codename atlas", "s1"),
                ("atlas launch planning", "s1"),
                ("project codename atlas", "s2"),
            ],
            "project codename atlas",
        ),
        (
            "hybrid",
            "launch berlin timeline",
            [
                ("berlin office notes", "s1"),
                ("launch berlin timeline finalized", "s1"),
                ("launch checklist", "s2"),
            ],
            "launch berlin timeline finalized",
        ),
        (
            "semantic",
            "sparkly water preference",
            [
                ("coffee beans inventory", "s1"),
                ("sparkling water is preferred", "s1"),
                ("sparkling water is preferred", "s2"),
            ],
            "sparkling water is preferred",
        ),
    ],
)
async def test_retrieval_ranking_parity_between_backends(
    tmp_path,
    retrieval,
    query,
    entries,
    expected,
):
    in_memory_top = await _remember_top_content(
        InMemoryBackend(),
        retrieval=retrieval,
        query=query,
        entries=entries,
    )
    sqlite_top = await _remember_top_content(
        SQLiteBackend(str(tmp_path / f"{retrieval}.db")),
        retrieval=retrieval,
        query=query,
        entries=entries,
    )

    assert in_memory_top == expected
    assert sqlite_top == expected
