# Copyright (c) DevSecCode. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root.

"""
Token counting and budget tracking for workflow runs.

Uses tiktoken for accurate counts when available, falls back to a
simple word-based heuristic (1 token ≈ 0.75 words) otherwise.

Usage::

    budget = TokenBudget(max_tokens=100_000)
    budget.consume(1500, source="ai-node-1")
    if budget.exceeded:
        ...
"""

from __future__ import annotations

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

logger = logging.getLogger(__name__)

# -- tiktoken import (optional dependency) -----------------------------------

_encoder: Any = None

try:
	import tiktoken
	_encoder = tiktoken.encoding_for_model('gpt-4')  # cl100k_base
except Exception:
	pass  # tiktoken not installed — use heuristic


def count_tokens(text: str) -> int:
	"""Estimate the number of tokens in *text*."""
	if not text:
		return 0
	if _encoder is not None:
		return len(_encoder.encode(text))
	# Heuristic: ~1 token per 4 characters (conservative)
	return max(1, len(text) // 4)


@dataclass
class TokenBudget:
	"""Track token consumption against a per-run budget."""

	max_tokens: int = 100_000
	warn_pct: float = 0.8  # fire warning at 80%

	_consumed: int = field(default=0, init=False)
	_breakdown: dict[str, int] = field(default_factory=dict, init=False)

	def consume(self, tokens: int, source: str = 'unknown') -> None:
		"""Record *tokens* consumed by *source* (node id or label)."""
		self._consumed += tokens
		self._breakdown[source] = self._breakdown.get(source, 0) + tokens

	@property
	def consumed(self) -> int:
		return self._consumed

	@property
	def remaining(self) -> int:
		return max(0, self.max_tokens - self._consumed)

	@property
	def exceeded(self) -> bool:
		return self._consumed >= self.max_tokens

	@property
	def warning(self) -> bool:
		return self._consumed >= int(self.max_tokens * self.warn_pct)

	@property
	def pct_used(self) -> float:
		if self.max_tokens <= 0:
			return 1.0
		return self._consumed / self.max_tokens

	def to_dict(self) -> dict[str, Any]:
		return {
			'max_tokens': self.max_tokens,
			'consumed': self._consumed,
			'remaining': self.remaining,
			'pct_used': round(self.pct_used, 4),
			'exceeded': self.exceeded,
			'breakdown': dict(self._breakdown),
		}
