"""Tests for P0 token optimization: token counter + context window."""
from __future__ import annotations

from logos_adk.token_counter import count_text, count_messages, estimate_tool_tokens
from logos_adk.types import Message
from logos_adk.config import ContextWindowConfig, AgentConfig


class TestTokenCounterHeuristic:
    def test_count_text_empty(self):
        assert count_text("", model="unknown-model") == 0

    def test_count_text_heuristic_basic(self):
        result = count_text("Hello world, this is a test.", model="unknown-model")
        assert 5 <= result <= 10

    def test_count_text_cyrillic(self):
        result = count_text("Привет мир, это тест.", model="unknown-model")
        assert 5 <= result <= 15

    def test_count_messages_sums_roles_and_content(self):
        msgs = [
            Message(role="system", content="You are helpful."),
            Message(role="user", content="Hello"),
            Message(role="assistant", content="Hi there!"),
        ]
        total = count_messages(msgs, model="unknown-model")
        assert total > 10

    def test_count_messages_empty_list(self):
        assert count_messages([], model="unknown-model") == 0

    def test_estimate_tool_tokens_none(self):
        assert estimate_tool_tokens(None, model="unknown-model") == 0

    def test_estimate_tool_tokens_basic(self):
        tools = [
            {"name": "search", "description": "Search the web", "parameters": {"type": "object", "properties": {"q": {"type": "string"}}}},
        ]
        result = estimate_tool_tokens(tools, model="unknown-model")
        assert result > 10

    def test_count_text_cache_returns_same_result(self):
        text = "Consistent result for caching"
        r1 = count_text(text, model="unknown-model")
        r2 = count_text(text, model="unknown-model")
        assert r1 == r2


class TestContextWindowConfig:
    def test_defaults(self):
        cfg = ContextWindowConfig()
        assert cfg.max_context_tokens == 0
        assert cfg.reserved_output_tokens == 4096
        assert cfg.strategy == "chunked_summary"
        assert cfg.chunk_size == 10
        assert cfg.summary_max_tokens == 300
        assert cfg.max_summary_chain == 5
        assert cfg.keep_recent_messages == 6
        assert cfg.summary_model is None
        assert cfg.summary_system_prompt is None

    def test_agent_config_has_context_window(self):
        ac = AgentConfig(name="test")
        assert isinstance(ac.context_window, ContextWindowConfig)
        assert ac.context_window.max_context_tokens == 0


import asyncio
from logos_adk.context_window import trim_context, async_trim_context
from logos_adk.state import AgentState


def _make_messages(n: int, *, role: str = "user", text_len: int = 100) -> list[Message]:
    """Create n messages with predictable content."""
    msgs = []
    for i in range(n):
        r = "user" if i % 2 == 0 else "assistant"
        if role != "user":
            r = role
        msgs.append(Message(role=r, content=f"Message {i}: " + "x" * text_len))
    return msgs


def _run(coro):
    return asyncio.run(coro)


class TestTruncateStrategy:
    def test_noop_when_under_limit(self):
        cfg = ContextWindowConfig(max_context_tokens=100_000, strategy="truncate", keep_recent_messages=4)
        msgs = [Message(role="system", content="Be helpful.")] + _make_messages(4)
        state = AgentState()
        result = trim_context(msgs, config=cfg, state=state, model="unknown")
        assert len(result) == 5  # unchanged

    def test_truncate_removes_oldest_keeps_system_and_recent(self):
        cfg = ContextWindowConfig(max_context_tokens=200, strategy="truncate", keep_recent_messages=2)
        system = Message(role="system", content="System prompt.")
        old = _make_messages(10, text_len=50)
        msgs = [system] + old
        state = AgentState()
        result = trim_context(msgs, config=cfg, state=state, model="unknown")
        # Must keep system + last 2 messages at minimum
        assert result[0].role == "system"
        assert result[0].content == "System prompt."
        assert len(result) <= len(msgs)
        assert result[-1].content == old[-1].content
        assert result[-2].content == old[-2].content

    def test_truncate_disabled_when_max_zero(self):
        cfg = ContextWindowConfig(max_context_tokens=0, strategy="truncate")
        msgs = _make_messages(20, text_len=200)
        state = AgentState()
        result = trim_context(msgs, config=cfg, state=state, model="unknown")
        assert len(result) == 20


class TestSlidingWindowStrategy:
    def test_sliding_window_creates_summary_message(self):
        cfg = ContextWindowConfig(max_context_tokens=200, strategy="sliding_window", keep_recent_messages=2)
        system = Message(role="system", content="System.")
        old = _make_messages(10, text_len=50)
        msgs = [system] + old
        state = AgentState()

        # When no provider is available, sliding_window falls back to truncate
        result = trim_context(msgs, config=cfg, state=state, model="unknown")
        assert result[0].role == "system"
        assert len(result) < len(msgs)


class TestChunkedSummaryStrategy:
    def _make_summarize_fn(self, responses: list[str] | None = None):
        """Create a mock summarize_fn that returns predictable text."""
        call_count = [0]
        call_log: list[tuple[list[Message], str | None]] = []

        async def _summarize(msgs: list[Message], prompt: str | None) -> str:
            call_log.append((msgs, prompt))
            idx = call_count[0]
            call_count[0] += 1
            if responses and idx < len(responses):
                return responses[idx]
            return f"Summary of {len(msgs)} messages"

        _summarize.call_log = call_log  # type: ignore
        _summarize.call_count = call_count  # type: ignore
        return _summarize

    def test_noop_under_limit(self):
        cfg = ContextWindowConfig(max_context_tokens=100_000, strategy="chunked_summary", chunk_size=5)
        msgs = _make_messages(4)
        state = AgentState()
        result = _run(async_trim_context(msgs, config=cfg, state=state, model="unknown"))
        assert len(result) == 4
        assert "_context_summaries" not in state.metadata

    def test_chunk_creates_summary_when_buffer_full(self):
        cfg = ContextWindowConfig(
            max_context_tokens=300,
            strategy="chunked_summary",
            chunk_size=4,
            keep_recent_messages=2,
            max_summary_chain=5,
        )
        msgs = [Message(role="system", content="System.")] + _make_messages(12, text_len=50)
        state = AgentState()
        summarize_fn = self._make_summarize_fn()

        result = _run(async_trim_context(
            msgs, config=cfg, state=state, model="unknown", summarize_fn=summarize_fn,
        ))
        # Should have created summaries
        chain = state.metadata.get("_context_summaries", [])
        assert len(chain) > 0
        # Result must contain system + summary message + recent
        assert result[0].role == "system"
        assert result[0].content == "System."
        # At least one summary message should exist
        summary_msgs = [m for m in result if "[Conversation history summaries]" in m.content]
        assert len(summary_msgs) <= 1  # condensed into one

    def test_chain_persists_in_state_metadata(self):
        cfg = ContextWindowConfig(
            max_context_tokens=300,
            strategy="chunked_summary",
            chunk_size=3,
            keep_recent_messages=2,
            max_summary_chain=10,
        )
        msgs = [Message(role="system", content="S.")] + _make_messages(10, text_len=50)
        state = AgentState()
        summarize_fn = self._make_summarize_fn()

        _run(async_trim_context(msgs, config=cfg, state=state, model="unknown", summarize_fn=summarize_fn))

        assert "_context_summaries" in state.metadata
        assert isinstance(state.metadata["_context_summaries"], list)
        assert "_context_summarized_count" in state.metadata
        assert state.metadata["_context_summarized_count"] > 0

    def test_chain_merges_when_exceeds_max(self):
        cfg = ContextWindowConfig(
            max_context_tokens=200,
            strategy="chunked_summary",
            chunk_size=2,
            keep_recent_messages=2,
            max_summary_chain=2,
        )
        # Enough messages to produce >2 chunks
        msgs = [Message(role="system", content="S.")] + _make_messages(14, text_len=40)
        state = AgentState()
        summarize_fn = self._make_summarize_fn()

        _run(async_trim_context(msgs, config=cfg, state=state, model="unknown", summarize_fn=summarize_fn))

        chain = state.metadata.get("_context_summaries", [])
        assert len(chain) <= cfg.max_summary_chain

    def test_recent_messages_always_preserved(self):
        cfg = ContextWindowConfig(
            max_context_tokens=300,
            strategy="chunked_summary",
            chunk_size=3,
            keep_recent_messages=3,
        )
        msgs = _make_messages(12, text_len=50)
        state = AgentState()
        summarize_fn = self._make_summarize_fn()

        result = _run(async_trim_context(msgs, config=cfg, state=state, model="unknown", summarize_fn=summarize_fn))

        # Last 3 messages must be preserved exactly
        assert result[-1].content == msgs[-1].content
        assert result[-2].content == msgs[-2].content
        assert result[-3].content == msgs[-3].content

    def test_falls_back_to_truncate_without_summarize_fn(self):
        cfg = ContextWindowConfig(max_context_tokens=200, strategy="chunked_summary", keep_recent_messages=2)
        msgs = _make_messages(20, text_len=50)
        state = AgentState()

        # No summarize_fn — should fall back to truncate
        result = _run(async_trim_context(msgs, config=cfg, state=state, model="unknown", summarize_fn=None))
        assert len(result) < 20
        assert result[-1].content == msgs[-1].content


from logos_adk.agent import Agent


class TestAgentContextWindowBuilder:
    def test_context_window_builder_sets_config(self):
        agent = Agent("test").context_window(
            max_tokens=128_000,
            strategy="chunked_summary",
            chunk_size=8,
            keep_recent=4,
        )
        cfg = agent._config.context_window
        assert cfg.max_context_tokens == 128_000
        assert cfg.strategy == "chunked_summary"
        assert cfg.chunk_size == 8
        assert cfg.keep_recent_messages == 4

    def test_context_window_builder_defaults(self):
        agent = Agent("test").context_window(max_tokens=50_000)
        cfg = agent._config.context_window
        assert cfg.max_context_tokens == 50_000
        assert cfg.strategy == "chunked_summary"
        assert cfg.chunk_size == 10
        assert cfg.keep_recent_messages == 6

    def test_context_window_disabled_by_default(self):
        agent = Agent("test")
        assert agent._config.context_window.max_context_tokens == 0


class TestAgentIntegration:
    def test_build_provider_messages_trims_when_configured(self):
        agent = Agent("trim-test").context_window(max_tokens=200, strategy="truncate", keep_recent=2)

        conversation = _make_messages(20, text_len=50)
        memory_messages: list[Message] = []
        state = AgentState()

        result = agent._build_provider_messages(conversation, memory_messages, state=state)

        assert len(result) < 22
        assert result[-1].content == conversation[-1].content
        assert result[-2].content == conversation[-2].content

    def test_build_provider_messages_untouched_when_disabled(self):
        agent = Agent("no-trim")
        conversation = _make_messages(20, text_len=50)
        memory_messages: list[Message] = []

        result = agent._build_provider_messages(conversation, memory_messages)
        assert len(result) == 20

    def test_build_provider_messages_untouched_when_no_state(self):
        agent = Agent("no-state").context_window(max_tokens=200, strategy="truncate", keep_recent=2)
        conversation = _make_messages(20, text_len=50)
        result = agent._build_provider_messages(conversation, [])
        assert len(result) == 20


class TestContextMetrics:
    def test_count_messages_accounts_for_tool_results(self):
        from logos_adk.types import ToolCall, ToolResult
        msgs = [
            Message(role="user", content="Use the tool"),
            Message(
                role="assistant",
                content="Calling search",
                tool_calls=[ToolCall(id="tc1", name="search", arguments={"q": "test"})],
            ),
            Message(
                role="tool",
                content="Result here",
                tool_result=ToolResult(call_id="tc1", content="Search result: found 5 items"),
            ),
        ]
        total = count_messages(msgs, model="unknown")
        assert total > 20

    def test_estimate_tool_tokens_empty_list(self):
        assert estimate_tool_tokens([], model="unknown") == 0


class TestConfigParser:
    def test_parse_context_window_from_yaml_data(self):
        from logos_adk.config_parser import parse_agent_config_data
        data = {
            "name": "test-agent",
            "model": "claude-sonnet-4",
            "context_window": {
                "max_context_tokens": 100_000,
                "strategy": "chunked_summary",
                "chunk_size": 8,
                "keep_recent_messages": 4,
            },
        }
        cfg = parse_agent_config_data(data)
        assert cfg.context_window.max_context_tokens == 100_000
        assert cfg.context_window.strategy == "chunked_summary"
        assert cfg.context_window.chunk_size == 8
        assert cfg.context_window.keep_recent_messages == 4

    def test_parse_context_window_defaults_when_absent(self):
        from logos_adk.config_parser import parse_agent_config_data
        data = {"name": "test-agent", "model": "claude-sonnet-4"}
        cfg = parse_agent_config_data(data)
        assert cfg.context_window.max_context_tokens == 0

    def test_parse_context_window_rejects_unknown_fields(self):
        from logos_adk.config_parser import parse_agent_config_data
        from logos_adk.errors import ConfigError
        import pytest
        data = {
            "name": "test-agent",
            "context_window": {"bogus_field": 123},
        }
        with pytest.raises(ConfigError):
            parse_agent_config_data(data)


class TestPublicAPI:
    def test_context_window_config_importable_from_root(self):
        from logos_adk import ContextWindowConfig
        cfg = ContextWindowConfig()
        assert cfg.max_context_tokens == 0

    def test_token_counter_importable_from_root(self):
        from logos_adk.token_counter import count_text
        assert count_text("hello", model="") > 0
