"""Token budget tracker for memory injection decisions.

Provides token counting (exact via Anthropic API or heuristic fallback),
session budget tracking, and a should_inject_memories() gate that decides
whether injecting memories into context is worth the token cost.

Usage:
    tracker = TokenTracker()
    count = await tracker.count_tokens("some text")
    tracker.record_usage("session-123", input_tokens=500, output_tokens=200)
    should_inject = tracker.should_inject_memories(
        session_id="session-123",
        candidate_tokens=1500,
        context_budget=200_000,
    )
"""

from __future__ import annotations

import hashlib
import logging
import time
from dataclasses import dataclass, field
from typing import Any

logger = logging.getLogger(__name__)

# Default configuration
DEFAULT_CACHE_TTL_SECONDS = 3600  # 1 hour
DEFAULT_SESSION_TTL_SECONDS = 86400  # 24 hours
DEFAULT_INJECTION_THRESHOLD = 0.15  # Don't inject if would use > 15% of remaining budget
DEFAULT_MODEL = "claude-sonnet-4-20250514"

# Heuristic: average English chars per token (Anthropic-recommended)
CHARS_PER_TOKEN = 3.5


@dataclass
class CacheEntry:
    """A cached token count with expiration."""

    token_count: int
    expires_at: float  # time.monotonic() value


@dataclass
class SessionBudget:
    """Tracks token usage within a session."""

    session_id: str
    input_tokens: int = 0
    output_tokens: int = 0
    memory_tokens_injected: int = 0
    created_at: float = field(default_factory=time.monotonic)
    updated_at: float = field(default_factory=time.monotonic)

    @property
    def total_tokens(self) -> int:
        """Total tokens consumed in this session."""
        return self.input_tokens + self.output_tokens

    def record(
        self,
        input_tokens: int = 0,
        output_tokens: int = 0,
        memory_tokens: int = 0,
    ) -> None:
        """Record token usage."""
        self.input_tokens += input_tokens
        self.output_tokens += output_tokens
        self.memory_tokens_injected += memory_tokens
        self.updated_at = time.monotonic()


class TokenTracker:
    """Token budget tracker with caching and session management.

    Supports two counting modes:
        1. Exact: Uses Anthropic's count_tokens API (requires ANTHROPIC_API_KEY)
        2. Heuristic: chars/4 estimate (always available, no API calls)

    Token counts are cached by content hash (SHA-256) with configurable TTL.
    """

    def __init__(
        self,
        *,
        cache_ttl: int = DEFAULT_CACHE_TTL_SECONDS,
        session_ttl: int = DEFAULT_SESSION_TTL_SECONDS,
        model: str = DEFAULT_MODEL,
        use_api: bool = False,
    ):
        """Initialize the token tracker.

        Args:
            cache_ttl: How long to cache token counts (seconds). Default: 1 hour.
            session_ttl: How long to keep session budgets (seconds). Default: 24 hours.
            model: Anthropic model name for count_tokens API.
            use_api: If True, attempt to use Anthropic API for exact counting.
                     Falls back to heuristic if API unavailable.
        """
        self.cache_ttl = cache_ttl
        self.session_ttl = session_ttl
        self.model = model
        self.use_api = use_api

        self._cache: dict[str, CacheEntry] = {}
        self._sessions: dict[str, SessionBudget] = {}
        self._client: Any = None
        self._api_available: bool | None = None  # None = not yet checked

    def _content_hash(self, text: str) -> str:
        """Compute SHA-256 hash for cache key."""
        return hashlib.sha256(text.encode("utf-8")).hexdigest()

    def _heuristic_count(self, text: str) -> int:
        """Estimate token count using character heuristic.

        Args:
            text: Input text.

        Returns:
            Estimated token count (minimum 1).
        """
        return max(1, int(len(text) / CHARS_PER_TOKEN))

    async def _get_client(self) -> Any:
        """Lazy-load Anthropic async client."""
        if self._client is None:
            try:
                import os

                from anthropic import AsyncAnthropic

                # Only create client if API key is available
                api_key = os.environ.get("ANTHROPIC_API_KEY")
                if not api_key:
                    self._api_available = False
                    return None

                self._client = AsyncAnthropic(api_key=api_key)
                self._api_available = True
            except ImportError:
                logger.debug("anthropic package not installed, using heuristic counting")
                self._api_available = False
                return None
            except Exception as e:
                logger.warning("Failed to initialize Anthropic client: %s", e)
                self._api_available = False
                return None
        return self._client

    def _check_cache(self, content_hash: str) -> int | None:
        """Check cache for a token count.

        Returns:
            Cached token count, or None if not cached/expired.
        """
        entry = self._cache.get(content_hash)
        if entry is None:
            return None
        if time.monotonic() > entry.expires_at:
            del self._cache[content_hash]
            return None
        return entry.token_count

    def _store_cache(self, content_hash: str, token_count: int) -> None:
        """Store a token count in the cache."""
        self._cache[content_hash] = CacheEntry(
            token_count=token_count,
            expires_at=time.monotonic() + self.cache_ttl,
        )

    async def count_tokens(self, text: str) -> int:
        """Count tokens in text, using cache and API/heuristic.

        Priority:
            1. Check cache (by content hash)
            2. If use_api=True and API available: exact count via count_tokens
            3. Fallback: chars/4 heuristic

        Args:
            text: Text to count tokens for.

        Returns:
            Token count (exact or estimated).
        """
        content_hash = self._content_hash(text)

        # Check cache first
        cached = self._check_cache(content_hash)
        if cached is not None:
            return cached

        token_count: int

        # Try API if enabled
        if self.use_api and self._api_available is not False:
            client = await self._get_client()
            if client is not None:
                try:
                    result = await client.messages.count_tokens(
                        model=self.model,
                        messages=[{"role": "user", "content": text}],
                    )
                    token_count = result.input_tokens
                    self._store_cache(content_hash, token_count)
                    return token_count
                except Exception as e:
                    logger.warning("count_tokens API failed, using heuristic: %s", e)

        # Heuristic fallback
        token_count = self._heuristic_count(text)
        self._store_cache(content_hash, token_count)
        return token_count

    def count_tokens_sync(self, text: str) -> int:
        """Synchronous token counting (heuristic only, uses cache).

        For sync contexts where async API call isn't practical.

        Args:
            text: Text to count tokens for.

        Returns:
            Token count (from cache or heuristic).
        """
        content_hash = self._content_hash(text)
        cached = self._check_cache(content_hash)
        if cached is not None:
            return cached

        token_count = self._heuristic_count(text)
        self._store_cache(content_hash, token_count)
        return token_count

    # ------------------------------------------------------------------
    # Session budget management
    # ------------------------------------------------------------------

    def get_session(self, session_id: str) -> SessionBudget:
        """Get or create a session budget tracker.

        Args:
            session_id: Unique session identifier.

        Returns:
            SessionBudget for the session.
        """
        self._cleanup_expired_sessions()

        if session_id not in self._sessions:
            self._sessions[session_id] = SessionBudget(session_id=session_id)
        return self._sessions[session_id]

    def record_usage(
        self,
        session_id: str,
        input_tokens: int = 0,
        output_tokens: int = 0,
        memory_tokens: int = 0,
    ) -> SessionBudget:
        """Record token usage for a session.

        Args:
            session_id: Session identifier.
            input_tokens: Input tokens consumed.
            output_tokens: Output tokens generated.
            memory_tokens: Tokens from memory injection specifically.

        Returns:
            Updated SessionBudget.
        """
        session = self.get_session(session_id)
        session.record(input_tokens, output_tokens, memory_tokens)
        return session

    def should_inject_memories(
        self,
        session_id: str,
        candidate_tokens: int,
        context_budget: int,
        threshold: float = DEFAULT_INJECTION_THRESHOLD,
    ) -> bool:
        """Decide whether injecting memories is worth the token cost.

        Gate function: returns True if injecting `candidate_tokens` worth of
        memories would use less than `threshold` fraction of remaining budget.

        Args:
            session_id: Session identifier.
            candidate_tokens: How many tokens the candidate memories would use.
            context_budget: Total context window budget (e.g., 200_000).
            threshold: Max fraction of remaining budget to use. Default: 0.15 (15%).

        Returns:
            True if injection is recommended, False if budget is too tight.
        """
        session = self.get_session(session_id)
        remaining = max(0, context_budget - session.input_tokens)

        if remaining == 0:
            return False

        injection_fraction = candidate_tokens / remaining
        return injection_fraction <= threshold

    def _cleanup_expired_sessions(self) -> int:
        """Remove expired sessions.

        Returns:
            Number of sessions removed.
        """
        now = time.monotonic()
        expired = [
            sid for sid, s in self._sessions.items() if (now - s.created_at) > self.session_ttl
        ]
        for sid in expired:
            del self._sessions[sid]
        return len(expired)

    def cleanup_cache(self) -> int:
        """Remove expired cache entries.

        Returns:
            Number of entries removed.
        """
        now = time.monotonic()
        expired = [key for key, entry in self._cache.items() if now > entry.expires_at]
        for key in expired:
            del self._cache[key]
        return len(expired)

    def stats(self) -> dict[str, Any]:
        """Return tracker statistics.

        Returns:
            Dict with cache_size, active_sessions, api_available.
        """
        self._cleanup_expired_sessions()
        self.cleanup_cache()
        return {
            "cache_size": len(self._cache),
            "active_sessions": len(self._sessions),
            "api_available": self._api_available,
            "model": self.model,
            "use_api": self.use_api,
        }
