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

import pytest

from ag2 import Context, ToolResult
from ag2.events import (
    ModelRequest,
    ModelResponse,
    TextInput,
    ToolCallEvent,
    ToolCallsEvent,
    ToolResultEvent,
    ToolResultsEvent,
)
from ag2.policies.sliding_window import SlidingWindowPolicy


def _tool_response(call_id: str = "tc_1", name: str = "get") -> ModelResponse:
    return ModelResponse(
        message=None,
        tool_calls=ToolCallsEvent(calls=[ToolCallEvent(id=call_id, name=name, arguments="{}")]),
    )


def _tool_results(parent_id: str = "tc_1", name: str = "get") -> ToolResultsEvent:
    return ToolResultsEvent(
        results=[ToolResultEvent(parent_id=parent_id, name=name, result=ToolResult("ok"))],
    )


class TestNoTrimming:
    @pytest.mark.asyncio
    async def test_events_within_limit_are_unchanged(self, context: Context) -> None:
        events = [ModelRequest([TextInput("a")]), ModelRequest([TextInput("b")])]
        policy = SlidingWindowPolicy(max_events=5)

        prompts, result = await policy.apply([], events, context)

        assert result == events
        assert prompts == []

    @pytest.mark.asyncio
    async def test_events_at_exact_limit(self, context: Context) -> None:
        events = [ModelRequest([TextInput("a")]), ModelRequest([TextInput("b")])]
        policy = SlidingWindowPolicy(max_events=2)

        _, result = await policy.apply([], events, context)

        assert result == events


class TestTrimming:
    @pytest.mark.asyncio
    async def test_keeps_last_n_events(self, context: Context) -> None:
        events = [ModelRequest([TextInput(str(i))]) for i in range(5)]
        policy = SlidingWindowPolicy(max_events=2)

        _, result = await policy.apply([], events, context)

        assert len(result) == 2
        assert result[0].parts[0].content == "3"
        assert result[1].parts[0].content == "4"

    @pytest.mark.asyncio
    async def test_transparent_adds_prompt(self, context: Context) -> None:
        events = [ModelRequest([TextInput(str(i))]) for i in range(5)]
        policy = SlidingWindowPolicy(max_events=2, transparent=True)

        prompts, result = await policy.apply(["existing"], events, context)

        assert len(result) == 2
        assert len(prompts) == 2
        assert prompts[0] == "existing"
        assert "2" in prompts[1] and "5" in prompts[1]


class TestOrphanedToolResults:
    """After trimming, leading ToolResultsEvents without a matching tool_use should be dropped."""

    @pytest.mark.asyncio
    async def test_leading_orphaned_tool_result_is_skipped(self, context: Context) -> None:
        events = [
            _tool_response("tc_1"),  # will be trimmed
            _tool_results("tc_1"),  # orphaned after trim — should be skipped
            ModelRequest([TextInput("next")]),
            ModelRequest([TextInput("final")]),
        ]
        policy = SlidingWindowPolicy(max_events=3)

        _, result = await policy.apply([], events, context)

        # The ToolResultsEvent should be dropped, leaving 2 events
        assert len(result) == 2
        assert isinstance(result[0], ModelRequest)
        assert result[0].parts[0].content == "next"

    @pytest.mark.asyncio
    async def test_multiple_leading_orphaned_tool_results_are_skipped(self, context: Context) -> None:
        events = [
            _tool_response("tc_1"),
            _tool_response("tc_2"),
            _tool_results("tc_1"),
            _tool_results("tc_2"),
            ModelRequest([TextInput("hello")]),
        ]
        policy = SlidingWindowPolicy(max_events=3)

        _, result = await policy.apply([], events, context)

        # Both orphaned ToolResultsEvents should be dropped
        assert len(result) == 1
        assert isinstance(result[0], ModelRequest)

    @pytest.mark.asyncio
    async def test_non_leading_tool_result_is_kept(self, context: Context) -> None:
        """ToolResultsEvent after a non-ToolResultsEvent should be preserved."""
        events = [
            ModelRequest([TextInput("old")]),
            _tool_response("tc_1"),
            _tool_results("tc_1"),
            ModelRequest([TextInput("new")]),
        ]
        policy = SlidingWindowPolicy(max_events=3)

        _, result = await policy.apply([], events, context)

        # trim keeps last 3: [tool_response, tool_results, request]
        assert len(result) == 3
        assert isinstance(result[0], ModelResponse)
        assert isinstance(result[1], ToolResultsEvent)
        assert isinstance(result[2], ModelRequest)

    @pytest.mark.asyncio
    async def test_transparent_count_reflects_skipped_orphans(self, context: Context) -> None:
        events = [
            _tool_response("tc_1"),
            _tool_results("tc_1"),
            ModelRequest([TextInput("a")]),
            ModelRequest([TextInput("b")]),
        ]
        policy = SlidingWindowPolicy(max_events=3, transparent=True)

        prompts, result = await policy.apply([], events, context)

        assert len(result) == 2
        # Prompt should reflect actual count after orphan removal
        assert "2" in prompts[-1] and "4" in prompts[-1]

    @pytest.mark.asyncio
    async def test_mid_window_orphaned_tool_result_is_dropped(self, context: Context) -> None:
        """An orphaned ToolResultsEvent at any position must be dropped, not only at the head."""
        events = [
            _tool_response("tc_1"),  # trimmed
            ModelRequest([TextInput("a")]),
            _tool_results("tc_1"),  # orphan inside window (parent trimmed)
            ModelRequest([TextInput("b")]),
        ]
        policy = SlidingWindowPolicy(max_events=3)

        _, result = await policy.apply([], events, context)

        assert len(result) == 2
        assert all(not isinstance(e, ToolResultsEvent) for e in result)
        assert [e.parts[0].content for e in result] == ["a", "b"]

    @pytest.mark.asyncio
    async def test_paired_tool_call_and_result_within_window_are_kept(self, context: Context) -> None:
        """When the matching ToolCallsEvent survives the window, its results stay."""
        events = [
            ModelRequest([TextInput("old")]),  # trimmed
            _tool_response("tc_1"),
            _tool_results("tc_1"),
            ModelRequest([TextInput("done")]),
        ]
        policy = SlidingWindowPolicy(max_events=3)

        _, result = await policy.apply([], events, context)

        assert len(result) == 3
        assert isinstance(result[1], ToolResultsEvent)

    @pytest.mark.asyncio
    async def test_orphan_results_at_multiple_positions_are_dropped(self, context: Context) -> None:
        """Orphans at head, middle, and elsewhere must all be removed in one pass."""
        events = [
            _tool_response("tc_1"),  # trimmed
            _tool_response("tc_2"),  # trimmed
            _tool_response("tc_3"),  # trimmed
            _tool_results("tc_1"),  # orphan
            ModelRequest([TextInput("a")]),
            _tool_results("tc_2"),  # orphan
            ModelRequest([TextInput("b")]),
            _tool_results("tc_3"),  # orphan
        ]
        policy = SlidingWindowPolicy(max_events=5)

        _, result = await policy.apply([], events, context)

        assert all(not isinstance(e, ToolResultsEvent) for e in result)
        assert [e.parts[0].content for e in result] == ["a", "b"]