"""Token estimation utilities for memory flush threshold calculations."""

from __future__ import annotations

import json
from typing import Any


def estimate_tokens(text: str) -> int:
    """Estimate token count for a string.

    Tries tiktoken (cl100k_base) first; falls back to len(text) // 4.
    """
    try:
        import tiktoken

        enc = tiktoken.get_encoding("cl100k_base")
        return len(enc.encode(text))
    except Exception:
        return max(1, len(text) // 4)


def estimate_messages_tokens(messages: list[dict[str, Any]]) -> int:
    """Sum token estimates across a list of session messages."""
    total = 0
    for m in messages:
        content = m.get("content")
        if isinstance(content, str):
            total += estimate_tokens(content)
        elif isinstance(content, list):
            for block in content:
                if isinstance(block, dict):
                    total += estimate_tokens(block.get("text", ""))
        # Count tool call arguments
        for tc in m.get("tool_calls", []):
            args = tc.get("function", {}).get("arguments", "")
            if isinstance(args, str):
                total += estimate_tokens(args)
            elif isinstance(args, dict):
                total += estimate_tokens(json.dumps(args, ensure_ascii=False))
    return total
