#!/usr/bin/env python3
"""
Functional tests for Context Window Management and Token Counter.
Tests pre-flight processing and context trimming strategies.
"""

import asyncio
import json
import os
import sys

sys.path.insert(0, os.path.dirname(__file__))


# ============================================================================
# TEST 1: Token Counter - Basic counting
# ============================================================================


def test_token_counter_basic():
    """Test basic token counting."""
    from logos_adk.token_counter import count_text, count_messages
    from logos_adk.types import Message

    # Test text counting
    text = "Hello, world! This is a test."
    tokens = count_text(text, model="gpt-4")
    print(f"  Text: '{text}'")
    print(f"  Tokens: {tokens}")
    assert tokens > 0

    # Test message counting
    messages = [
        Message(role="user", content="Hello"),
        Message(role="assistant", content="Hi there!"),
        Message(role="user", content="How are you?"),
    ]
    total = count_messages(messages, model="gpt-4")
    print(f"  3 messages total tokens: {total}")
    assert total > 0

    print("✓ test_token_counter_basic PASSED")
    return True


def test_token_counter_heuristic():
    """Test heuristic counting fallback."""
    from logos_adk.token_counter import count_text, _heuristic_count, _is_mostly_latin

    # Test latin detection
    assert _is_mostly_latin("Hello world") == True
    assert _is_mostly_latin("Привет мир") == False

    # Test heuristic
    latin_text = "Hello world"  # 11 chars / 4 = ~3 tokens
    cyrillic_text = "Привет мир"  # 9 chars / 2 = ~5 tokens

    latin_tokens = _heuristic_count(latin_text)
    cyrillic_tokens = _heuristic_count(cyrillic_text)

    print(f"  Latin '{latin_text}': {latin_tokens} tokens")
    print(f"  Cyrillic '{cyrillic_text}': {cyrillic_tokens} tokens")

    assert latin_tokens >= 2
    assert cyrillic_tokens >= 3

    print("✓ test_token_counter_heuristic PASSED")
    return True


def test_token_counter_with_tools():
    """Test tool token estimation."""
    from logos_adk.token_counter import estimate_tool_tokens

    tools = [
        {
            "type": "function",
            "function": {
                "name": "get_weather",
                "description": "Get weather for location",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "location": {"type": "string", "description": "City name"},
                        "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
                    },
                    "required": ["location"],
                },
            },
        }
    ]

    tokens = estimate_tool_tokens(tools, model="gpt-4")
    print(f"  Tool schema tokens: {tokens}")
    assert tokens > 0

    # Empty tools
    assert estimate_tool_tokens(None) == 0
    assert estimate_tool_tokens([]) == 0

    print("✓ test_token_counter_with_tools PASSED")
    return True


# ============================================================================
# TEST 2: Context Window - Truncation strategy
# ============================================================================


def test_context_truncation():
    """Test context truncation strategy."""
    from logos_adk.context_window import trim_context, _truncate
    from logos_adk.config import ContextWindowConfig
    from logos_adk.state import AgentState
    from logos_adk.types import Message

    # Create many messages that exceed the limit
    messages = [Message(role="user", content=f"Message number {i}") for i in range(20)]

    config = ContextWindowConfig(
        max_context_tokens=100,
        strategy="truncate",
        keep_recent_messages=3,
    )

    state = AgentState()

    # Trim with a very small limit
    trimmed = trim_context(
        messages,
        config=config,
        state=state,
        model="gpt-4",
    )

    print(f"  Original messages: {len(messages)}")
    print(f"  Trimmed messages: {len(trimmed)}")
    print(f"  Config: keep_recent={config.keep_recent_messages}")

    # Should keep recent messages
    assert len(trimmed) <= len(messages)

    print("✓ test_context_truncation PASSED")
    return True


def test_context_preserves_system():
    """Test that system messages are preserved during trimming."""
    from logos_adk.context_window import _truncate
    from logos_adk.types import Message

    messages = [
        Message(role="system", content="You are a helpful assistant."),
        Message(role="user", content="Hello"),
        Message(role="assistant", content="Hi!"),
        Message(role="user", content="How are you?"),
        Message(role="assistant", content="Good, thanks!"),
    ]

    # Very restrictive limit
    trimmed = _truncate(messages, max_tokens=10, keep_recent=1, model="gpt-4")

    # System message should always be preserved
    roles = [m.role for m in trimmed]
    print(f"  Roles after trimming: {roles}")

    assert "system" in roles, "System message should be preserved"

    print("✓ test_context_preserves_system PASSED")
    return True


# ============================================================================
# TEST 3: Context Window - Split system and conversation
# ============================================================================


def test_split_system_conversation():
    """Test splitting system messages from conversation."""
    from logos_adk.context_window import _split_system_and_conversation
    from logos_adk.types import Message

    messages = [
        Message(role="system", content="System prompt 1"),
        Message(role="system", content="System prompt 2"),
        Message(role="user", content="Hello"),
        Message(role="assistant", content="Hi"),
        Message(role="user", content="How are you?"),
    ]

    system, conv = _split_system_and_conversation(messages)

    print(f"  System messages: {len(system)}")
    print(f"  Conversation messages: {len(conv)}")

    assert len(system) == 2
    assert len(conv) == 3
    assert all(m.role == "system" for m in system)
    assert all(m.role != "system" for m in conv)

    # No system messages
    messages_no_system = [Message(role="user", content="Hello")]
    system2, conv2 = _split_system_and_conversation(messages_no_system)
    assert len(system2) == 0
    assert len(conv2) == 1

    print("✓ test_split_system_conversation PASSED")
    return True


# ============================================================================
# TEST 4: Context Window - Configuration
# ============================================================================


def test_context_window_config():
    """Test ContextWindowConfig."""
    from logos_adk.config import ContextWindowConfig

    # Default config
    config = ContextWindowConfig()
    assert config.max_context_tokens == 0  # 0 = disabled
    assert config.strategy == "chunked_summary"
    assert config.keep_recent_messages == 6
    assert config.reserved_output_tokens == 4096

    # Custom config
    config = ContextWindowConfig(
        max_context_tokens=4000,
        strategy="sliding_window",
        keep_recent_messages=5,
        reserved_output_tokens=500,
    )

    assert config.max_context_tokens == 4000
    assert config.strategy == "sliding_window"
    assert config.keep_recent_messages == 5
    assert config.reserved_output_tokens == 500

    print("✓ test_context_window_config PASSED")
    return True


# ============================================================================
# TEST 5: Context Window - Integration with Agent
# ============================================================================


async def test_agent_context_trimming():
    """Test that agent properly trims context."""
    from logos_adk.agent import Agent
    from logos_adk.providers.openai import OpenAICompatibleProvider

    # Reset runtime
    from logos_adk.runtime import _default_runtime_holder

    _default_runtime_holder.reset()

    # Create agent with small context window
    provider = OpenAICompatibleProvider(
        model=os.getenv("OPENAI_COMPAT_MODEL", "gpt-4o-mini"),
        base_url=os.getenv("OPENAI_COMPAT_BASE_URL", "https://api.example.com/v1"),
        api_key=os.getenv("OPENAI_COMPAT_API_KEY", "test-api-key"),
    )

    agent = Agent(
        "test_agent",
        model=provider,
    )

    # Run agent to verify it works with context management
    response = await agent.run("Say 'hello' in one word.")

    print(f"  Response: {response.content[:50]}...")
    assert len(response.content) > 0

    print("✓ test_agent_context_trimming PASSED")
    return True


# ============================================================================
# TEST 6: Real API - Test with actual token counting
# ============================================================================


async def test_real_api_token_counting():
    """Test with real API - verify token counting works."""
    from logos_adk.token_counter import count_messages, count_text
    from logos_adk.agent import Agent
    from logos_adk.providers.openai import OpenAICompatibleProvider
    from logos_adk.types import Message

    from logos_adk.runtime import _default_runtime_holder

    _default_runtime_holder.reset()

    provider = OpenAICompatibleProvider(
        model=os.getenv("OPENAI_COMPAT_MODEL", "gpt-4o-mini"),
        base_url=os.getenv("OPENAI_COMPAT_BASE_URL", "https://api.example.com/v1"),
        api_key=os.getenv("OPENAI_COMPAT_API_KEY", "test-api-key"),
    )

    # Count tokens for typical conversation
    messages = [
        Message(role="system", content="You are a helpful assistant."),
        Message(role="user", content="What is quantum computing?"),
        Message(role="assistant", content="Quantum computing is a type of computation..."),
        Message(role="user", content="Explain more about qubits."),
        Message(role="assistant", content="A qubit (quantum bit) is the basic unit..."),
    ]

    # Count with different models
    gpt4_tokens = count_messages(messages, model="gpt-4")
    claude_tokens = count_messages(messages, model="claude-3")
    default_tokens = count_messages(messages, model="")

    print(f"  GPT-4 tokens: {gpt4_tokens}")
    print(f"  Claude tokens: {claude_tokens}")
    print(f"  Default (heuristic) tokens: {default_tokens}")

    # All should be > 0 and reasonable
    assert gpt4_tokens > 50
    assert claude_tokens > 50
    assert default_tokens > 50

    print("✓ test_real_api_token_counting PASSED")
    return True


# ============================================================================
# TEST 7: Context - Sliding Window Strategy
# ============================================================================


async def test_sliding_window_strategy():
    """Test sliding window summarization strategy."""
    from logos_adk.context_window import async_trim_context
    from logos_adk.config import ContextWindowConfig
    from logos_adk.state import AgentState
    from logos_adk.types import Message

    # Create long conversation
    messages = [
        Message(role="system", content="You are a helpful assistant."),
    ]
    for i in range(10):
        messages.append(Message(role="user", content=f"User message {i}"))
        messages.append(Message(role="assistant", content=f"Assistant response {i}"))

    config = ContextWindowConfig(
        max_context_tokens=200,
        strategy="sliding_window",
        keep_recent_messages=2,
    )
    state = AgentState()

    # Summarize function (mock)
    async def mock_summarize(msgs, prompt):
        return f"Summarized {len(msgs)} messages into a brief summary."

    # Trim with sliding window (will use summarize_fn)
    trimmed = await async_trim_context(
        messages,
        config=config,
        state=state,
        model="gpt-4",
        summarize_fn=mock_summarize,
    )

    print(f"  Original: {len(messages)} messages")
    print(f"  After sliding window: {len(trimmed)} messages")

    # Should have fewer messages (old ones summarized)
    assert len(trimmed) <= len(messages)

    print("✓ test_sliding_window_strategy PASSED")
    return True


# ============================================================================
# MAIN
# ============================================================================


async def main():
    print("=" * 60)
    print("CONTEXT WINDOW & TOKEN COUNTER TESTS")
    print("=" * 60)

    results = []

    # Sync tests
    sync_tests = [
        ("Token Counter: Basic", test_token_counter_basic),
        ("Token Counter: Heuristic", test_token_counter_heuristic),
        ("Token Counter: Tools", test_token_counter_with_tools),
        ("Context: Truncation", test_context_truncation),
        ("Context: Preserve System", test_context_preserves_system),
        ("Context: Split System/Conv", test_split_system_conversation),
        ("Context: Config", test_context_window_config),
    ]

    for name, test_func in sync_tests:
        try:
            result = test_func()
            results.append((name, result, None))
        except Exception as e:
            print(f"✗ {name}: {e}")
            import traceback

            traceback.print_exc()
            results.append((name, False, str(e)))

    # Async tests
    async_tests = [
        ("Agent: Context Trimming", test_agent_context_trimming),
        ("Token Counter: Real API", test_real_api_token_counting),
        ("Context: Sliding Window", test_sliding_window_strategy),
    ]

    for name, test_func in async_tests:
        try:
            result = await test_func()
            results.append((name, result, None))
        except Exception as e:
            print(f"✗ {name}: {e}")
            import traceback

            traceback.print_exc()
            results.append((name, False, str(e)))

    print("\n" + "=" * 60)
    print("RESULTS")
    print("=" * 60)

    passed = sum(1 for _, ok, _ in results if ok)
    total = len(results)

    for name, ok, err in results:
        status = "✓ PASS" if ok else "✗ FAIL"
        print(f"{status}: {name}")
        if err:
            print(f"    Error: {err}")

    print(f"\n{passed}/{total} tests passed")
    return passed == total


if __name__ == "__main__":
    success = asyncio.run(main())
    sys.exit(0 if success else 1)
