# Copyright (c) 2026, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
#
# SPDX-License-Identifier: Apache-2.0

"""Tests for assembly policies (replacement for old ContextHarness tests)."""

import pytest

from ag2 import Context
from ag2.assembly import AssemblerMiddleware
from ag2.compact import CompactionSummary
from ag2.events import (
    ModelMessage,
    ModelReasoning,
    ModelRequest,
    ModelResponse,
    ObserverAlert,
    ProviderReplay,
    Severity,
    TextInput,
    ToolCallEvent,
    ToolResultEvent,
)
from ag2.knowledge import KnowledgeStore, MemoryKnowledgeStore
from ag2.policies import (
    ConversationPolicy,
    EpisodicMemoryPolicy,
    SlidingWindowPolicy,
    TokenBudgetPolicy,
    WorkingMemoryPolicy,
)
from ag2.stream import MemoryStream


class ProviderReasoning(ModelReasoning, ProviderReplay):
    """Stands in for OpenAIReasoningEvent, keeping these tests provider-free."""

    __transient__ = False
    __replay_role__ = "anchor"


class TestConversationPolicy:
    @pytest.mark.asyncio
    async def test_filters_to_conversation_events(self) -> None:
        policy = ConversationPolicy()
        events = [
            ModelRequest([TextInput("hello")]),
            ModelResponse(message=ModelMessage(content="hi")),
            ToolCallEvent(name="search", arguments="{}"),
            ToolResultEvent(id="1", name="search", content="result"),
            ObserverAlert(source="mon", severity=Severity.WARNING, message="warn"),
        ]
        ctx = Context(stream=MemoryStream())
        prompts, filtered = await policy.apply([], events, ctx)
        assert len(filtered) == 4
        assert all(not isinstance(e, ObserverAlert) for e in filtered)

    @pytest.mark.asyncio
    async def test_drops_streamed_thinking(self) -> None:
        policy = ConversationPolicy()
        request = ModelRequest([TextInput("hello")])
        events = [ModelReasoning("thinking out loud"), request]
        ctx = Context(stream=MemoryStream())
        _, filtered = await policy.apply([], events, ctx)
        assert filtered == [request]

    @pytest.mark.asyncio
    async def test_keeps_provider_replay_with_tool_call(self) -> None:
        policy = ConversationPolicy()
        events = [
            ModelRequest([TextInput("hello")]),
            ProviderReasoning("planning the search"),
            ToolCallEvent(name="search", arguments="{}"),
            ToolResultEvent(id="1", name="search", content="result"),
        ]
        ctx = Context(stream=MemoryStream())
        _, filtered = await policy.apply([], events, ctx)
        assert filtered == events

    @pytest.mark.asyncio
    async def test_includes_compaction_summary(self) -> None:
        policy = ConversationPolicy()
        summary = CompactionSummary(summary="Earlier context...", event_count=50)
        events = [summary, ModelRequest([TextInput("hello")])]
        ctx = Context(stream=MemoryStream())
        _, filtered = await policy.apply([], events, ctx)
        assert summary in filtered


# NOTE: TestNetworkPolicy (V2 NetworkPolicy + FormattedEvent wrapping of
# DelegationResult / SchedulerTriggerFired / TopicMessage) was removed with
# the V2 rewrite. V3 will reintroduce equivalent assembly glue in Phase 2 /
# Phase 4 when network Tasks and the multi-participant session types land.


class TestSlidingWindowPolicy:
    @pytest.mark.asyncio
    async def test_no_trim_below_max(self) -> None:
        policy = SlidingWindowPolicy(max_events=10)
        events = [ModelRequest([TextInput(f"msg-{i}")]) for i in range(5)]
        ctx = Context(stream=MemoryStream())
        prompts, filtered = await policy.apply([], events, ctx)
        assert len(filtered) == 5
        assert prompts == []

    @pytest.mark.asyncio
    async def test_trims_to_max(self) -> None:
        policy = SlidingWindowPolicy(max_events=3)
        events = [ModelRequest([TextInput(f"msg-{i}")]) for i in range(10)]
        ctx = Context(stream=MemoryStream())
        _, filtered = await policy.apply([], events, ctx)
        assert len(filtered) == 3
        assert filtered[0].parts[0].content == "msg-7"

    @pytest.mark.asyncio
    async def test_transparent_adds_note(self) -> None:
        policy = SlidingWindowPolicy(max_events=3, transparent=True)
        events = [ModelRequest([TextInput(f"msg-{i}")]) for i in range(10)]
        ctx = Context(stream=MemoryStream())
        prompts, _ = await policy.apply([], events, ctx)
        assert len(prompts) == 1
        assert "3 of 10" in prompts[0]


class TestTokenBudgetPolicy:
    @pytest.mark.asyncio
    async def test_no_trim_within_budget(self) -> None:
        policy = TokenBudgetPolicy(max_tokens=10000)
        events = [ModelRequest([TextInput("short")])]
        ctx = Context(stream=MemoryStream())
        _, filtered = await policy.apply([], events, ctx)
        assert len(filtered) == 1

    @pytest.mark.asyncio
    async def test_trims_to_budget(self) -> None:
        policy = TokenBudgetPolicy(max_tokens=10, chars_per_token=1)
        events = [ModelRequest([TextInput("a" * 20)]), ModelRequest([TextInput("b" * 5)])]
        ctx = Context(stream=MemoryStream())
        _, filtered = await policy.apply([], events, ctx)
        # Should keep at least the last event that fits
        assert len(filtered) >= 1
        assert filtered[-1].parts[0].content == "b" * 5


class TestAssemblerMiddleware:
    @pytest.mark.asyncio
    async def test_applies_policies_in_order(self) -> None:
        stream = MemoryStream()
        ctx = Context(stream=stream)
        initial_event = ModelRequest([TextInput("start")])

        policy = ConversationPolicy()
        mw = AssemblerMiddleware(initial_event, ctx, policies=[policy])

        all_events = [
            ModelRequest([TextInput("hello")]),
            ModelResponse(message=ModelMessage(content="hi")),
            ObserverAlert(source="mon", severity=Severity.WARNING, message="warn"),
        ]

        received_events = None

        async def mock_llm_call(events, context):
            nonlocal received_events
            received_events = list(events)
            return ModelResponse(message=ModelMessage(content="response"))

        await mw.on_llm_call(mock_llm_call, all_events, ctx)

        assert received_events is not None
        assert len(received_events) == 2
        assert all(not isinstance(e, ObserverAlert) for e in received_events)

    @pytest.mark.asyncio
    async def test_restores_prompts_after_call(self) -> None:
        stream = MemoryStream()
        ctx = Context(stream=stream, prompt=["original"])
        initial_event = ModelRequest([TextInput("start")])

        class _PromptAdder:
            name = "adder"

            async def apply(self, prompts, events, context):
                return prompts + ["injected"], events

        mw = AssemblerMiddleware(initial_event, ctx, policies=[_PromptAdder()])

        async def mock_llm_call(events, context):
            assert "injected" in context.prompt
            return ModelResponse(message=ModelMessage(content="ok"))

        await mw.on_llm_call(mock_llm_call, [ModelRequest([TextInput("hi")])], ctx)
        assert ctx.prompt == ["original"]

    def test_validate_order_warns_on_bad_ordering(self) -> None:
        class _FakePolicy:
            def __init__(self, n):
                self.name = n

            async def apply(self, p, e, c):
                return p, e

        policies = [_FakePolicy("sliding_window"), _FakePolicy("episodic_memory")]
        warnings = AssemblerMiddleware.validate_order(policies)
        assert len(warnings) == 1
        assert "sliding_window" in warnings[0]

    def test_validate_order_no_warnings_for_correct_order(self) -> None:
        class _FakePolicy:
            def __init__(self, n):
                self.name = n

            async def apply(self, p, e, c):
                return p, e

        policies = [_FakePolicy("episodic_memory"), _FakePolicy("sliding_window")]
        warnings = AssemblerMiddleware.validate_order(policies)
        assert len(warnings) == 0

    @pytest.mark.asyncio
    async def test_restores_prompts_on_exception(self) -> None:
        stream = MemoryStream()
        ctx = Context(stream=stream, prompt=["original"])
        initial_event = ModelRequest([TextInput("start")])

        class _PromptAdder:
            name = "adder"

            async def apply(self, prompts, events, context):
                return prompts + ["injected"], events

        mw = AssemblerMiddleware(initial_event, ctx, policies=[_PromptAdder()])

        async def failing_llm_call(events, context):
            raise RuntimeError("LLM failed")

        with pytest.raises(RuntimeError):
            await mw.on_llm_call(failing_llm_call, [ModelRequest([TextInput("hi")])], ctx)

        # Prompts must be restored even after exception
        assert ctx.prompt == ["original"]


class TestEpisodicMemoryPolicy:
    @pytest.mark.asyncio
    async def test_injects_summaries_from_store(self) -> None:
        store = MemoryKnowledgeStore()
        await store.write("/memory/conversations/20260101T120000_abc.md", "Summary of session 1.")
        await store.write("/memory/conversations/20260102T120000_def.md", "Summary of session 2.")

        policy = EpisodicMemoryPolicy(max_episodes=5)
        ctx = Context(stream=MemoryStream())
        ctx.dependencies[KnowledgeStore] = store

        prompts, events = await policy.apply([], [ModelRequest([TextInput("hi")])], ctx)
        assert any("Past Conversations" in p for p in prompts)
        assert any("Summary of session 1" in p for p in prompts)
        assert any("Summary of session 2" in p for p in prompts)

    @pytest.mark.asyncio
    async def test_limits_to_max_episodes(self) -> None:
        store = MemoryKnowledgeStore()
        for i in range(10):
            await store.write(f"/memory/conversations/2026010{i}T120000_s{i}.md", f"Summary {i}")

        policy = EpisodicMemoryPolicy(max_episodes=3)
        ctx = Context(stream=MemoryStream())
        ctx.dependencies[KnowledgeStore] = store

        prompts, _ = await policy.apply([], [ModelRequest([TextInput("hi")])], ctx)
        # Should have the last 3 (most recent by sorted name)
        combined = " ".join(prompts)
        assert "Summary 7" in combined
        assert "Summary 8" in combined
        assert "Summary 9" in combined
        assert "Summary 0" not in combined

    @pytest.mark.asyncio
    async def test_no_op_without_store(self) -> None:
        policy = EpisodicMemoryPolicy()
        ctx = Context(stream=MemoryStream())
        prompts, events = await policy.apply(["existing"], [ModelRequest([TextInput("hi")])], ctx)
        assert prompts == ["existing"]

    @pytest.mark.asyncio
    async def test_no_op_when_no_summaries(self) -> None:
        store = MemoryKnowledgeStore()
        policy = EpisodicMemoryPolicy()
        ctx = Context(stream=MemoryStream())
        ctx.dependencies[KnowledgeStore] = store
        prompts, _ = await policy.apply([], [ModelRequest([TextInput("hi")])], ctx)
        assert prompts == []


class TestWorkingMemoryPolicy:
    @pytest.mark.asyncio
    async def test_injects_working_memory(self) -> None:
        store = MemoryKnowledgeStore()
        await store.write("/memory/working.md", "Current state: working on project X.")

        policy = WorkingMemoryPolicy()
        ctx = Context(stream=MemoryStream())
        ctx.dependencies[KnowledgeStore] = store

        prompts, _ = await policy.apply([], [ModelRequest([TextInput("hi")])], ctx)
        assert any("Working Memory" in p for p in prompts)
        assert any("project X" in p for p in prompts)

    @pytest.mark.asyncio
    async def test_no_op_without_store(self) -> None:
        policy = WorkingMemoryPolicy()
        ctx = Context(stream=MemoryStream())
        prompts, _ = await policy.apply([], [ModelRequest([TextInput("hi")])], ctx)
        assert prompts == []

    @pytest.mark.asyncio
    async def test_no_op_without_working_memory_file(self) -> None:
        store = MemoryKnowledgeStore()
        policy = WorkingMemoryPolicy()
        ctx = Context(stream=MemoryStream())
        ctx.dependencies[KnowledgeStore] = store
        prompts, _ = await policy.apply([], [ModelRequest([TextInput("hi")])], ctx)
        assert prompts == []


# NOTE: TestTopicInboxPolicy removed — TopicInboxPolicy was V2-network-specific.
# Equivalent SessionInboxPolicy coverage lives with the V3 network tests.