"""Tests for P1 token optimizations."""
from __future__ import annotations

from logos_adk.agent import Agent
from logos_adk.tools import tool


@tool
async def dummy_search(query: str) -> str:
    """Search for something."""
    return f"result for {query}"


class TestToolSchemaCache:
    def test_build_tool_schemas_returns_same_object_on_repeated_calls(self):
        agent = Agent("test").tools([dummy_search])
        first = agent._build_tool_schemas()
        second = agent._build_tool_schemas()
        assert first is second

    def test_tool_schema_cache_clears_on_invalidation(self):
        agent = Agent("test").tools([dummy_search])
        first = agent._build_tool_schemas()
        agent._invalidate_tool_schema_cache()
        second = agent._build_tool_schemas()
        assert first is not second
        assert first == second

    def test_tool_schema_cache_none_when_no_tools(self):
        agent = Agent("test")
        assert agent._build_tool_schemas() is None


class TestSystemPromptCache:
    def test_resolved_prompt_cached_within_run_scope(self):
        agent = Agent("test")
        agent._config.system_prompt = "Be helpful."
        agent._begin_run_caches()
        first = agent._resolved_system_prompt()
        second = agent._resolved_system_prompt()
        assert first is second
        agent._end_run_caches()

    def test_resolved_prompt_not_cached_outside_run_scope(self):
        agent = Agent("test")
        agent._config.system_prompt = "Be helpful."
        first = agent._resolved_system_prompt()
        second = agent._resolved_system_prompt()
        assert first == second

    def test_run_caches_clear_on_end(self):
        agent = Agent("test")
        agent._config.system_prompt = "Be helpful."
        agent._begin_run_caches()
        first = agent._resolved_system_prompt()
        agent._end_run_caches()
        agent._begin_run_caches()
        second = agent._resolved_system_prompt()
        agent._end_run_caches()
        assert first == second

    def test_run_cache_depth_supports_nesting(self):
        agent = Agent("test")
        agent._config.system_prompt = "Nested."
        agent._begin_run_caches()  # depth=1
        agent._begin_run_caches()  # depth=2
        first = agent._resolved_system_prompt()
        agent._end_run_caches()  # depth=1 — should NOT clear
        second = agent._resolved_system_prompt()
        assert first is second  # still cached
        agent._end_run_caches()  # depth=0 — clears


from logos_adk.team import Team
from logos_adk.crew import CrewTask


class TestCrewContextCompression:
    def test_build_task_prompt_truncates_long_context(self):
        team = Team("test-team")
        task = CrewTask(
            name="summarize",
            description="Summarize the research",
            depends_on=["research"],
            context=[],
        )
        long_output = "x" * 5000
        context = {"research": long_output}
        prompt = team._build_task_prompt(
            task,
            kickoff="Go",
            context=context,
            assigned_role=None,
            manager_instruction=None,
        )
        # Context should be truncated, not the full 5000 chars
        assert len(prompt) < 4000
        assert "[truncated]" in prompt

    def test_build_task_prompt_keeps_short_context(self):
        team = Team("test-team")
        task = CrewTask(
            name="use_result",
            description="Use the result",
            depends_on=["prior"],
            context=[],
        )
        short_output = "The answer is 42."
        context = {"prior": short_output}
        prompt = team._build_task_prompt(
            task,
            kickoff="Go",
            context=context,
            assigned_role=None,
            manager_instruction=None,
        )
        assert short_output in prompt
        assert "[truncated]" not in prompt


from logos_adk.config import ContextWindowConfig
from logos_adk.token_counter import count_text


class FakeMemoryEntry:
    """Minimal stand-in for MemoryEntry with just .content."""
    def __init__(self, content: str):
        self.content = content


class TestMemoryTokenBudget:
    def test_memory_max_tokens_default_is_zero(self):
        cfg = ContextWindowConfig()
        assert cfg.memory_max_tokens == 0

    def test_memory_messages_respects_token_budget(self):
        agent = Agent("test").context_window(max_tokens=10_000)
        agent._config.context_window.memory_max_tokens = 50
        entries = [FakeMemoryEntry("x" * 200) for _ in range(10)]
        msgs = agent._memory_messages(entries)
        if msgs:
            total_tokens = count_text(msgs[0].content, model="")
            assert total_tokens <= 70  # budget + overhead tolerance

    def test_memory_messages_preserves_all_when_under_budget(self):
        agent = Agent("test").context_window(max_tokens=10_000)
        agent._config.context_window.memory_max_tokens = 10_000
        entries = [FakeMemoryEntry("short") for _ in range(3)]
        msgs = agent._memory_messages(entries)
        assert len(msgs) == 1
        assert msgs[0].content.count("short") == 3

    def test_memory_messages_no_budget_means_no_limit(self):
        agent = Agent("test")
        entries = [FakeMemoryEntry("x" * 500) for _ in range(20)]
        msgs = agent._memory_messages(entries)
        assert len(msgs) == 1
        assert msgs[0].content.count("x" * 500) == 20
