"""Tests for adaptive memory learning."""
from __future__ import annotations

import pytest

from logos_adk.learning import MemoryLearningEngine
from logos_adk.memory import LongTermMemory, MemoryEntry, SessionMemory
from logos_adk.memory.backends import InMemoryBackend
from logos_adk.state import AgentState
from logos_adk.types import Message, Response, Usage


@pytest.mark.asyncio
async def test_memory_learning_reranks_entries_using_feedback_signals():
    backend = InMemoryBackend()
    engine = MemoryLearningEngine()

    preferred = engine.initialize_entry(MemoryEntry(content="Berlin rollout launch checklist", session_id="s1"))
    degraded = engine.initialize_entry(MemoryEntry(content="Berlin rollout launch checklist old", session_id="s1"))
    await backend.store(preferred)
    await backend.store(degraded)

    await engine.apply_feedback(backend, [preferred], helpful=True, reason="accepted_answer")
    await engine.apply_feedback(backend, [degraded], helpful=False, reason="wrong_memory")

    memory = LongTermMemory(backend=backend, retrieval="hybrid", limit=1, learning=engine)
    state = AgentState(session_id="s1")
    results = await memory.remember(
        [Message(role="user", content="launch checklist for berlin rollout")],
        state,
        session_id="s1",
    )

    assert len(results) == 1
    assert results[0].id == preferred.id
    assert state.metadata["adk_memory_entry_ids"] == [preferred.id]


@pytest.mark.asyncio
async def test_memory_learning_forgetting_policy_archives_low_value_entries():
    backend = InMemoryBackend()
    engine = MemoryLearningEngine(max_active_entries=2, archive_threshold=0.35, delete_threshold=0.1)

    low_value = engine.initialize_entry(MemoryEntry(content="stale project rumor", session_id="s1"))
    steady = engine.initialize_entry(MemoryEntry(content="current project status", session_id="s1"))
    strong = engine.initialize_entry(MemoryEntry(content="latest launch decision", session_id="s1"))
    for entry in (low_value, steady, strong):
        await backend.store(entry)

    for _ in range(3):
        await engine.apply_feedback(backend, [low_value], helpful=False, reason="bad_memory")
    await engine.apply_feedback(backend, [steady], helpful=True, reason="good_memory")
    await engine.apply_feedback(backend, [strong], helpful=True, reason="good_memory")
    await engine.enforce_forgetting_policy(backend, session_id="s1")

    remembered = await LongTermMemory(backend=backend, limit=5, learning=engine).remember(
        [Message(role="user", content="project launch")],
        AgentState(session_id="s1"),
        session_id="s1",
    )

    assert all(entry.id != low_value.id for entry in remembered)
    raw_entries = await backend.search("", session_id="s1", limit=10)
    archived = next(entry for entry in raw_entries if entry.id == low_value.id)
    assert archived.is_archived is True


@pytest.mark.asyncio
async def test_session_memory_uses_adaptive_summary_settings():
    engine = MemoryLearningEngine()
    memory = SessionMemory(strategy="summarize", max_messages=2, learning=engine)
    state = AgentState(
        session_id="s1",
        messages=[
            Message(role="user", content=f"fact {index}")
            if index % 2 == 0
            else Message(role="assistant", content=f"reply {index}")
            for index in range(16)
        ],
    )

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

    assert state.metadata["adk_memory_summary_strategy"] in {"balanced", "compact"}
    assert state.messages[0].content.startswith("Session summary:\n")
    assert len(state.messages[0].content.splitlines()) <= 15


@pytest.mark.asyncio
async def test_memory_learning_adjusts_context_budget():
    backend = InMemoryBackend()
    engine = MemoryLearningEngine(max_context_limit=5)
    entries = []
    for index in range(4):
        entry = engine.initialize_entry(MemoryEntry(content=f"project signal {index}", session_id="s1"))
        await backend.store(entry)
        entries.append(entry)
    for _ in range(3):
        for entry in entries:
            await engine.apply_feedback(backend, [entry], helpful=False, reason="weak_signal")

    memory = LongTermMemory(backend=backend, retrieval="keyword", limit=2, learning=engine)
    state = AgentState(session_id="s1")
    results = await memory.remember(
        [Message(role="user", content="project signal")],
        state,
        session_id="s1",
    )

    assert state.metadata["adk_memory_context_limit"] == 4
    assert len(results) == 4
