"""
Tests for context window truncation integration.

Tests the TruncationConfig, its integration with KivotRunner,
and the plugin configuration.
"""

import pytest
from datetime import timedelta
from unittest.mock import AsyncMock, patch

from .overrides import (
    ActivityParameters,
    KivotAgentsConfig,
    TruncationConfig,
    get_kivot_config,
    set_kivot_agent_overrides,
)
from .plugin import OpenAIAgentsKivotPlugin
from .runner import _maybe_truncate


# --- TruncationConfig tests ---


class TestTruncationConfig:
    """Test TruncationConfig defaults and behavior."""

    def test_defaults(self):
        config = TruncationConfig()
        assert config.max_messages == 100
        assert config.max_tokens == 0
        assert config.summarize_removed is False
        assert config.summary_model_ref == ""
        assert config.summary_max_tokens == 500
        assert config.preserve_system_messages is True

    def test_custom_values(self):
        config = TruncationConfig(
            max_messages=50,
            max_tokens=100000,
            summarize_removed=True,
            summary_model_ref="gpt-4o-mini",
            summary_max_tokens=300,
            preserve_system_messages=False,
        )
        assert config.max_messages == 50
        assert config.max_tokens == 100000
        assert config.summarize_removed is True
        assert config.summary_model_ref == "gpt-4o-mini"
        assert config.summary_max_tokens == 300
        assert config.preserve_system_messages is False


# --- KivotAgentsConfig with truncation ---


class TestConfigWithTruncation:
    """Test KivotAgentsConfig truncation field."""

    def test_config_without_truncation(self):
        config = KivotAgentsConfig(
            model_params=ActivityParameters(),
            tool_params=ActivityParameters(),
        )
        assert config.truncation is None

    def test_config_with_truncation(self):
        truncation = TruncationConfig(max_messages=50)
        config = KivotAgentsConfig(
            model_params=ActivityParameters(),
            tool_params=ActivityParameters(),
            truncation=truncation,
        )
        assert config.truncation is not None
        assert config.truncation.max_messages == 50

    def test_overrides_preserve_truncation(self):
        truncation = TruncationConfig(max_messages=30, summarize_removed=True)
        config = KivotAgentsConfig(
            model_params=ActivityParameters(),
            tool_params=ActivityParameters(),
            truncation=truncation,
        )
        with set_kivot_agent_overrides(config):
            current = get_kivot_config()
            assert current is not None
            assert current.truncation is not None
            assert current.truncation.max_messages == 30
            assert current.truncation.summarize_removed is True


# --- Plugin with truncation ---


class TestPluginWithTruncation:
    """Test OpenAIAgentsKivotPlugin truncation support."""

    def test_plugin_without_truncation(self):
        plugin = OpenAIAgentsKivotPlugin()
        assert plugin.config.truncation is None

    def test_plugin_with_truncation(self):
        truncation = TruncationConfig(
            max_messages=50,
            summarize_removed=True,
            summary_model_ref="gpt-4o-mini",
        )
        plugin = OpenAIAgentsKivotPlugin(truncation=truncation)
        assert plugin.config.truncation is not None
        assert plugin.config.truncation.max_messages == 50
        assert plugin.config.truncation.summarize_removed is True
        assert plugin.config.truncation.summary_model_ref == "gpt-4o-mini"


# --- _maybe_truncate function tests ---


class TestMaybeTruncate:
    """Test the _maybe_truncate helper function."""

    @pytest.mark.asyncio
    @patch("kivot_agents.runner.workflow")
    async def test_no_config_returns_messages_unchanged(self, mock_workflow):
        messages = [
            {"role": "system", "content": "You are helpful."},
            {"role": "user", "content": "Hello"},
        ]
        result = await _maybe_truncate(
            messages,
            truncation_config=None,
            run_id="run-1",
            thread_id=None,
            task_queue="core-activities",
            logger=mock_workflow.logger,
        )
        assert result is messages  # Same object, not copied

    @pytest.mark.asyncio
    @patch("kivot_agents.runner.workflow")
    async def test_below_threshold_no_activity_call(self, mock_workflow):
        messages = [
            {"role": "system", "content": "You are helpful."},
            {"role": "user", "content": "Hello"},
            {"role": "assistant", "content": "Hi!"},
        ]
        config = TruncationConfig(max_messages=10)
        result = await _maybe_truncate(
            messages,
            truncation_config=config,
            run_id="run-1",
            thread_id=None,
            task_queue="core-activities",
            logger=mock_workflow.logger,
        )
        # Should not call execute_activity
        mock_workflow.execute_activity.assert_not_called()
        assert result is messages

    @pytest.mark.asyncio
    @patch("kivot_agents.runner.workflow")
    async def test_above_threshold_calls_activity(self, mock_workflow):
        # Build messages exceeding the threshold
        messages = [{"role": "system", "content": "You are helpful."}]
        for i in range(20):
            messages.append({"role": "user", "content": f"msg {i}"})
            messages.append({"role": "assistant", "content": f"reply {i}"})
        # 1 system + 40 non-system = 41 messages

        config = TruncationConfig(max_messages=10)

        # Mock the activity response
        truncated_messages = messages[:1] + messages[-10:]  # system + last 10

        async def mock_execute(*args, **kwargs):
            return {
                "messages": truncated_messages,
                "removed_count": 30,
                "summary_inserted": False,
                "estimated_tokens": 100,
            }

        mock_workflow.execute_activity = AsyncMock(side_effect=mock_execute)

        result = await _maybe_truncate(
            messages,
            truncation_config=config,
            run_id="run-1",
            thread_id="thread-1",
            task_queue="core-activities",
            logger=mock_workflow.logger,
        )

        mock_workflow.execute_activity.assert_called_once()
        call_args = mock_workflow.execute_activity.call_args
        assert call_args[0][0] == "TruncateThreadActivity"

        # Verify the strategy was passed correctly
        activity_input = call_args[0][1]
        assert activity_input["strategy"]["max_messages"] == 10
        assert activity_input["strategy"]["preserve_system_messages"] is True
        assert activity_input["thread_id"] == "thread-1"

        assert len(result) == 11  # system + 10

    @pytest.mark.asyncio
    @patch("kivot_agents.runner.workflow")
    async def test_with_summary(self, mock_workflow):
        messages = [{"role": "system", "content": "inst"}]
        for i in range(20):
            messages.append({"role": "user", "content": f"q{i}"})

        config = TruncationConfig(
            max_messages=5,
            summarize_removed=True,
            summary_model_ref="gpt-4o-mini",
        )

        summary_msg = {
            "role": "system",
            "content": "[Previous conversation summary] ...",
            "tool_call_id": "__summary__",
        }
        truncated = [messages[0], summary_msg] + messages[-5:]

        async def mock_execute(*args, **kwargs):
            return {
                "messages": truncated,
                "removed_count": 15,
                "summary_inserted": True,
                "estimated_tokens": 80,
            }

        mock_workflow.execute_activity = AsyncMock(side_effect=mock_execute)

        result = await _maybe_truncate(
            messages,
            truncation_config=config,
            run_id="run-1",
            thread_id=None,
            task_queue="core-activities",
            logger=mock_workflow.logger,
        )

        assert len(result) == 7  # system + summary + 5 messages
        # Verify summarize_removed was passed
        activity_input = mock_workflow.execute_activity.call_args[0][1]
        assert activity_input["strategy"]["summarize_removed"] is True
        assert activity_input["strategy"]["summary_model_ref"] == "gpt-4o-mini"

    @pytest.mark.asyncio
    @patch("kivot_agents.runner.workflow")
    async def test_no_removal_returns_original(self, mock_workflow):
        """When activity says nothing was removed, return original messages."""
        messages = [
            {"role": "system", "content": "inst"},
            {"role": "user", "content": "hello"},
        ]
        # max_tokens > 0 triggers the activity call even with few messages
        config = TruncationConfig(max_messages=0, max_tokens=100000)

        async def mock_execute(*args, **kwargs):
            return {
                "messages": messages,
                "removed_count": 0,
                "summary_inserted": False,
                "estimated_tokens": 20,
            }

        mock_workflow.execute_activity = AsyncMock(side_effect=mock_execute)

        result = await _maybe_truncate(
            messages,
            truncation_config=config,
            run_id="run-1",
            thread_id=None,
            task_queue="core-activities",
            logger=mock_workflow.logger,
        )

        assert result is messages  # Original returned when no removal

    @pytest.mark.asyncio
    @patch("kivot_agents.runner.workflow")
    async def test_zero_limits_skips(self, mock_workflow):
        """When both max_messages and max_tokens are 0, skip truncation."""
        messages = [{"role": "user", "content": "hello"}] * 100
        config = TruncationConfig(max_messages=0, max_tokens=0)

        result = await _maybe_truncate(
            messages,
            truncation_config=config,
            run_id="run-1",
            thread_id=None,
            task_queue="core-activities",
            logger=mock_workflow.logger,
        )

        mock_workflow.execute_activity.assert_not_called()
        assert result is messages


if __name__ == "__main__":
    pytest.main([__file__, "-v"])
