"""Tests for TokenTracker module."""

from __future__ import annotations

import time
from unittest.mock import AsyncMock, patch

import pytest

from claude_engram.token_tracker import (
    CHARS_PER_TOKEN,
    DEFAULT_CACHE_TTL_SECONDS,
    DEFAULT_INJECTION_THRESHOLD,
    DEFAULT_SESSION_TTL_SECONDS,
    SessionBudget,
    TokenTracker,
)


class TestHeuristicCounting:
    """Tests for character-based token estimation."""

    def test_basic_count(self) -> None:
        tracker = TokenTracker()
        # 100 chars / 3.5 = 28 tokens
        text = "a" * 100
        assert tracker.count_tokens_sync(text) == 28

    def test_empty_string_returns_one(self) -> None:
        tracker = TokenTracker()
        assert tracker.count_tokens_sync("") == 1

    def test_short_string_returns_one(self) -> None:
        tracker = TokenTracker()
        assert tracker.count_tokens_sync("hi") == 1

    def test_exact_multiple(self) -> None:
        tracker = TokenTracker()
        # 3.5 * 10 = 35 chars -> int(35 / 3.5) = 10
        text = "a" * int(CHARS_PER_TOKEN * 10)
        assert tracker.count_tokens_sync(text) == 10

    def test_rounding_down(self) -> None:
        tracker = TokenTracker()
        # 5 chars / 3.5 = 1.43 -> int(1.43) = 1
        assert tracker.count_tokens_sync("hello") == 1


class TestCaching:
    """Tests for token count cache."""

    def test_cache_hit(self) -> None:
        tracker = TokenTracker()
        text = "test content for caching"
        count1 = tracker.count_tokens_sync(text)
        count2 = tracker.count_tokens_sync(text)
        assert count1 == count2
        assert tracker.stats()["cache_size"] == 1

    def test_different_texts_different_entries(self) -> None:
        tracker = TokenTracker()
        tracker.count_tokens_sync("text one")
        tracker.count_tokens_sync("text two")
        assert tracker.stats()["cache_size"] == 2

    def test_cache_expiry(self) -> None:
        tracker = TokenTracker(cache_ttl=1)  # 1 second TTL
        tracker.count_tokens_sync("ephemeral")
        assert tracker.stats()["cache_size"] == 1

        # Manually expire
        for entry in tracker._cache.values():
            entry.expires_at = time.monotonic() - 1

        tracker.cleanup_cache()
        assert tracker.stats()["cache_size"] == 0

    def test_cleanup_removes_expired(self) -> None:
        tracker = TokenTracker(cache_ttl=1)
        tracker.count_tokens_sync("item1")
        tracker.count_tokens_sync("item2")
        tracker.count_tokens_sync("item3")

        # Expire all
        for entry in tracker._cache.values():
            entry.expires_at = time.monotonic() - 1

        removed = tracker.cleanup_cache()
        assert removed == 3
        assert len(tracker._cache) == 0


class TestAsyncCounting:
    """Tests for async count_tokens method."""

    @pytest.mark.asyncio()
    async def test_async_heuristic_fallback(self) -> None:
        """When use_api=False, falls back to heuristic."""
        tracker = TokenTracker(use_api=False)
        count = await tracker.count_tokens("a" * 100)
        assert count == 28  # 100 / 3.5 = 28

    @pytest.mark.asyncio()
    async def test_async_uses_cache(self) -> None:
        tracker = TokenTracker(use_api=False)
        count1 = await tracker.count_tokens("cached text")
        count2 = await tracker.count_tokens("cached text")
        assert count1 == count2
        assert tracker.stats()["cache_size"] == 1

    @pytest.mark.asyncio()
    async def test_api_disabled_by_default(self) -> None:
        tracker = TokenTracker()
        assert tracker.use_api is False
        count = await tracker.count_tokens("test")
        assert count == 1  # 4 chars / 3.5 = 1

    @pytest.mark.asyncio()
    async def test_api_failure_falls_back_to_heuristic(self) -> None:
        """When API call fails, gracefully falls back to heuristic."""
        tracker = TokenTracker(use_api=True)
        tracker._api_available = True

        mock_client = AsyncMock()
        mock_client.messages.count_tokens.side_effect = Exception("API error")
        tracker._client = mock_client

        count = await tracker.count_tokens("a" * 35)
        assert count == 10  # heuristic: 35/3.5

    @pytest.mark.asyncio()
    async def test_no_api_key_disables_api(self) -> None:
        """Without ANTHROPIC_API_KEY, api_available is set to False."""
        tracker = TokenTracker(use_api=True)
        with patch.dict("os.environ", {}, clear=True):
            client = await tracker._get_client()
            assert client is None
            assert tracker._api_available is False


class TestSessionBudget:
    """Tests for SessionBudget dataclass."""

    def test_initial_values(self) -> None:
        budget = SessionBudget(session_id="test-1")
        assert budget.input_tokens == 0
        assert budget.output_tokens == 0
        assert budget.memory_tokens_injected == 0
        assert budget.total_tokens == 0

    def test_record_usage(self) -> None:
        budget = SessionBudget(session_id="test-1")
        budget.record(input_tokens=100, output_tokens=50, memory_tokens=30)
        assert budget.input_tokens == 100
        assert budget.output_tokens == 50
        assert budget.memory_tokens_injected == 30
        assert budget.total_tokens == 150

    def test_cumulative_recording(self) -> None:
        budget = SessionBudget(session_id="test-1")
        budget.record(input_tokens=100)
        budget.record(input_tokens=200, output_tokens=50)
        assert budget.input_tokens == 300
        assert budget.output_tokens == 50
        assert budget.total_tokens == 350


class TestSessionManagement:
    """Tests for session lifecycle in TokenTracker."""

    def test_get_or_create_session(self) -> None:
        tracker = TokenTracker()
        session = tracker.get_session("s1")
        assert session.session_id == "s1"
        assert session.total_tokens == 0

    def test_same_session_returned(self) -> None:
        tracker = TokenTracker()
        s1 = tracker.get_session("s1")
        s2 = tracker.get_session("s1")
        assert s1 is s2

    def test_different_sessions_independent(self) -> None:
        tracker = TokenTracker()
        s1 = tracker.get_session("s1")
        s2 = tracker.get_session("s2")
        s1.record(input_tokens=100)
        assert s2.input_tokens == 0

    def test_record_usage_via_tracker(self) -> None:
        tracker = TokenTracker()
        session = tracker.record_usage("s1", input_tokens=500, output_tokens=200)
        assert session.total_tokens == 700

    def test_session_expiry(self) -> None:
        tracker = TokenTracker(session_ttl=1)
        tracker.get_session("s1")
        assert tracker.stats()["active_sessions"] == 1

        # Manually expire
        for s in tracker._sessions.values():
            s.created_at = time.monotonic() - 2

        tracker._cleanup_expired_sessions()
        assert tracker.stats()["active_sessions"] == 0


class TestShouldInjectMemories:
    """Tests for the memory injection gate function."""

    def test_inject_when_budget_available(self) -> None:
        tracker = TokenTracker()
        # 1000 tokens into 200K budget = 0.5% << 15% threshold
        assert tracker.should_inject_memories(
            session_id="s1",
            candidate_tokens=1000,
            context_budget=200_000,
        )

    def test_reject_when_budget_tight(self) -> None:
        tracker = TokenTracker()
        # Use up most of the budget
        tracker.record_usage("s1", input_tokens=195_000)
        # 1000 tokens into remaining 5K = 20% > 15% threshold
        assert not tracker.should_inject_memories(
            session_id="s1",
            candidate_tokens=1000,
            context_budget=200_000,
        )

    def test_reject_when_zero_remaining(self) -> None:
        tracker = TokenTracker()
        tracker.record_usage("s1", input_tokens=200_000)
        assert not tracker.should_inject_memories(
            session_id="s1",
            candidate_tokens=100,
            context_budget=200_000,
        )

    def test_custom_threshold(self) -> None:
        tracker = TokenTracker()
        tracker.record_usage("s1", input_tokens=150_000)
        # remaining = 50K, 10K injection = 20%
        # Default threshold (15%) -> reject
        assert not tracker.should_inject_memories(
            session_id="s1",
            candidate_tokens=10_000,
            context_budget=200_000,
            threshold=0.15,
        )
        # Higher threshold (25%) -> accept
        assert tracker.should_inject_memories(
            session_id="s1",
            candidate_tokens=10_000,
            context_budget=200_000,
            threshold=0.25,
        )

    def test_default_threshold_value(self) -> None:
        assert DEFAULT_INJECTION_THRESHOLD == 0.15

    def test_small_injection_always_passes(self) -> None:
        tracker = TokenTracker()
        # 10 tokens / 200K = 0.005% - always fine
        assert tracker.should_inject_memories(
            session_id="s1",
            candidate_tokens=10,
            context_budget=200_000,
        )


class TestStats:
    """Tests for tracker statistics."""

    def test_initial_stats(self) -> None:
        tracker = TokenTracker()
        stats = tracker.stats()
        assert stats["cache_size"] == 0
        assert stats["active_sessions"] == 0
        assert stats["api_available"] is None  # Not yet checked
        assert stats["use_api"] is False

    def test_stats_after_usage(self) -> None:
        tracker = TokenTracker()
        tracker.count_tokens_sync("text1")
        tracker.count_tokens_sync("text2")
        tracker.get_session("s1")
        tracker.get_session("s2")
        stats = tracker.stats()
        assert stats["cache_size"] == 2
        assert stats["active_sessions"] == 2


class TestConstants:
    """Tests for module constants."""

    def test_cache_ttl_one_hour(self) -> None:
        assert DEFAULT_CACHE_TTL_SECONDS == 3600

    def test_session_ttl_24_hours(self) -> None:
        assert DEFAULT_SESSION_TTL_SECONDS == 86400

    def test_chars_per_token(self) -> None:
        assert CHARS_PER_TOKEN == 3.5

    def test_injection_threshold(self) -> None:
        assert DEFAULT_INJECTION_THRESHOLD == 0.15
