"""Token counting with a hard accuracy rule: tiktoken is NEVER a stand-in for Claude.

Why this module is strict about it
----------------------------------
tiktoken implements OpenAI's BPE vocabularies (cl100k_base, o200k_base). Claude
models use a different, proprietary tokenizer. Anthropic's own documentation
notes that tiktoken undercounts Claude tokens by 15-20%+ on normal English text
(and worse on code). On top of that, Claude models from Opus 4.7 onward use a
newer tokenizer that produces roughly 30% more tokens for the same text than
earlier Claude models — so even a "calibrated" tiktoken fudge factor is wrong
across Claude generations. This is a documented finding from prior research on
this project, not a style preference.

Consequences baked into this module's design:

* :func:`count_tokens_tiktoken` raises ``ValueError`` if given a ``claude-*``
  model name. There is no override.
* :func:`count_tokens` never silently falls back to tiktoken for a Claude
  model. If ``ANTHROPIC_API_KEY`` is unset, it raises ``RuntimeError`` telling
  the user to set the key — a missing number is better than a wrong one.
* The only trustworthy Claude count is Anthropic's own
  ``client.messages.count_tokens`` API (:func:`count_tokens_anthropic`).

Every :class:`TokenCount` records the ``method`` used ("anthropic_api" or
"tiktoken_approx") so downstream reporting can never confuse a real Claude
count with a GPT-family approximation.
"""

from __future__ import annotations

import os
from dataclasses import dataclass

try:  # pragma: no cover - optional convenience, not a hard dependency
    from dotenv import load_dotenv

    load_dotenv()  # loads a .env in the current/parent dir into os.environ, if present
except ImportError:
    pass

__all__ = [
    "TokenCount",
    "count_tokens",
    "count_tokens_anthropic",
    "count_tokens_tiktoken",
]


@dataclass
class TokenCount:
    tokens: int
    model: str
    method: str  # "anthropic_api" or "tiktoken_approx"


def count_tokens_anthropic(text: str, model: str = "claude-sonnet-5") -> TokenCount:
    """Count tokens with the real Anthropic count_tokens API.

    Uses ``client.messages.count_tokens`` — the only accurate source of Claude
    token counts (see module docstring for why tiktoken is never acceptable
    here). Requires the ``ANTHROPIC_API_KEY`` environment variable and the
    ``anthropic`` package.
    """
    try:
        import anthropic  # lazy: package must stay importable without the SDK
    except ImportError as exc:  # pragma: no cover - depends on environment
        raise RuntimeError(
            "The 'anthropic' package is required for Claude token counts: "
            "pip install anthropic"
        ) from exc

    client = anthropic.Anthropic()
    response = client.messages.count_tokens(
        model=model,
        messages=[{"role": "user", "content": text}],
    )
    return TokenCount(tokens=response.input_tokens, model=model, method="anthropic_api")


def count_tokens_tiktoken(text: str, model: str = "gpt-4o") -> TokenCount:
    """Approximate GPT-family token count via tiktoken.

    NEVER valid for Claude models: tiktoken undercounts Claude tokens by
    15-20%+ on normal text (more on code), and Claude tokenizers changed again
    at Opus 4.7+ (~30% more tokens for the same text). Passing a ``claude-*``
    model name here raises ``ValueError`` — use
    :func:`count_tokens_anthropic` instead.
    """
    if model.lower().startswith("claude"):
        raise ValueError(
            f"tiktoken must never be used for Claude models (got {model!r}): "
            "it undercounts Claude tokens by 15-20%+ and cannot track the "
            "Opus 4.7+ tokenizer change. Use count_tokens_anthropic()."
        )
    try:
        import tiktoken  # lazy: optional dependency
    except ImportError as exc:  # pragma: no cover - depends on environment
        raise RuntimeError(
            "The 'tiktoken' package is required for GPT-family token counts: "
            "pip install 'aave-compress[tiktoken]'"
        ) from exc

    try:
        encoding = tiktoken.encoding_for_model(model)
    except KeyError:
        encoding = tiktoken.get_encoding("o200k_base")
    return TokenCount(
        tokens=len(encoding.encode(text)), model=model, method="tiktoken_approx"
    )


def count_tokens(text: str, model: str = "claude-sonnet-5") -> TokenCount:
    """Dispatch to the correct counter for ``model``.

    * ``claude-*`` → the Anthropic count_tokens API. If ``ANTHROPIC_API_KEY``
      is unset this raises ``RuntimeError`` — it does NOT fall back to
      tiktoken, which would silently violate the accuracy requirement in the
      module docstring.
    * ``gpt-*`` (and other non-Claude names) → tiktoken approximation.
    """
    if model.lower().startswith("claude"):
        if not os.environ.get("ANTHROPIC_API_KEY"):
            raise RuntimeError(
                "ANTHROPIC_API_KEY is not set, so an accurate Claude token "
                f"count for {model!r} is unavailable. Set ANTHROPIC_API_KEY "
                "to use the real Anthropic count_tokens API. Refusing to fall "
                "back to tiktoken: it undercounts Claude tokens by 15-20%+ "
                "and would produce untrustworthy numbers."
            )
        return count_tokens_anthropic(text, model=model)
    return count_tokens_tiktoken(text, model=model)