"""Pre-flight token counting for context window management."""
from __future__ import annotations

import json
from functools import lru_cache

from logos_adk.types import Message, content_text, normalize_content_parts

_MSG_OVERHEAD = 4
_REQUEST_OVERHEAD = 3


def _is_mostly_latin(text: str) -> bool:
    if not text:
        return True
    alpha = [c for c in text if c.isalpha()]
    if not alpha:
        return True
    latin = sum(1 for c in alpha if c.isascii())
    return latin / len(alpha) > 0.6


def _heuristic_count(text: str) -> int:
    if not text:
        return 0
    chars_per_token = 4.0 if _is_mostly_latin(text) else 2.0
    return max(1, int(len(text) / chars_per_token + 0.5))


@lru_cache(maxsize=4)
def _get_tiktoken_encoding(model: str):
    try:
        import tiktoken
    except ImportError:
        return None
    if "claude" in model.lower():
        return tiktoken.get_encoding("cl100k_base")
    try:
        return tiktoken.encoding_for_model(model)
    except Exception:
        try:
            return tiktoken.get_encoding("cl100k_base")
        except Exception:
            return None


_text_cache: dict[tuple[str, str], int] = {}
_TEXT_CACHE_MAX = 2048


def count_text(text: str, *, model: str = "") -> int:
    if not text:
        return 0
    cache_key = (text, model)
    cached = _text_cache.get(cache_key)
    if cached is not None:
        return cached

    enc = _get_tiktoken_encoding(model)
    if enc is not None:
        try:
            result = len(enc.encode(text))
        except Exception:
            result = _heuristic_count(text)
    else:
        result = _heuristic_count(text)

    if len(_text_cache) >= _TEXT_CACHE_MAX:
        keys = list(_text_cache.keys())
        for k in keys[: len(keys) // 2]:
            _text_cache.pop(k, None)
    _text_cache[cache_key] = result
    return result


def count_messages(messages: list[Message], *, model: str = "") -> int:
    if not messages:
        return 0
    total = _REQUEST_OVERHEAD
    for msg in messages:
        text = content_text(msg.content, normalize_content_parts(msg.content_parts))
        total += _MSG_OVERHEAD + count_text(text, model=model)
        if msg.tool_calls:
            for tc in msg.tool_calls:
                total += count_text(tc.name, model=model)
                total += count_text(
                    json.dumps(tc.arguments) if isinstance(tc.arguments, dict) else str(tc.arguments or ""),
                    model=model,
                )
        if msg.tool_result is not None:
            total += count_text(str(msg.tool_result.content or ""), model=model)
    return total


def estimate_tool_tokens(tools: list[dict] | None, *, model: str = "") -> int:
    if not tools:
        return 0
    return count_text(json.dumps(tools), model=model)
