import asyncio
import logging
from dataclasses import dataclass, field
from typing import Optional, Sequence

logger = logging.getLogger(__name__)


@dataclass(slots=True)
class RetryPolicy:
    max_attempts: int = 5
    base_delay: float = 0.5
    max_delay: float = 30.0
    retry_on: Sequence[int] = field(default_factory=lambda: (429, 500, 502, 503, 529))

    def delay_for(self, attempt: int) -> float:
        return min(self.base_delay * (2 ** attempt), self.max_delay)


class TokenCounter:
    """Counts tokens for a batch of documents against a model endpoint."""

    def __init__(self, client, model: str, policy: Optional[RetryPolicy] = None):
        self._client = client
        self._model = model
        self._policy = policy or RetryPolicy()

    async def count(self, text: str) -> int:
        last_error: Optional[Exception] = None
        for attempt in range(self._policy.max_attempts):
            try:
                response = await self._client.messages.count_tokens(
                    model=self._model,
                    messages=[{"role": "user", "content": text}],
                )
                return int(response.input_tokens)
            except Exception as exc:
                status = getattr(exc, "status_code", None)
                if status not in self._policy.retry_on:
                    raise
                last_error = exc
                wait = self._policy.delay_for(attempt)
                logger.warning("retry %s after %.1fs (status=%s)", attempt, wait, status)
                await asyncio.sleep(wait)
        raise RuntimeError(f"count_tokens failed after retries: {last_error}")

    async def count_many(self, texts: Sequence[str], concurrency: int = 4) -> list[int]:
        semaphore = asyncio.Semaphore(concurrency)

        async def _one(item: str) -> int:
            async with semaphore:
                return await self.count(item)

        return list(await asyncio.gather(*(_one(t) for t in texts)))


def net_tokens(single: int, doubled: int) -> int:
    """Envelope overhead cancels out when you subtract the single from the doubled."""
    return doubled - single