"""
Token Counter — Estimates token count for CLAUDE.md content.

Uses a simple approximation (words/0.75) when tiktoken is not available,
or tiktoken when it is. Provides guidance on whether the file is too long.
"""

import re


# Approximate: 1 token ~= 0.75 words for English text
WORDS_PER_TOKEN = 0.75

# Thresholds
WARN_TOKENS = 1500    # Start warning
CRITICAL_TOKENS = 3000  # Likely wasting context


class TokenCounter:
    """Counts tokens in CLAUDE.md content."""

    def __init__(self):
        self._encoder = None
        self._try_load_tiktoken()

    def _try_load_tiktoken(self):
        """Try to load tiktoken for accurate counting."""
        try:
            import tiktoken
            # Use cl100k_base which is close to Claude's tokenizer
            self._encoder = tiktoken.get_encoding("cl100k_base")
        except ImportError:
            pass

    def count(self, text: str) -> int:
        """Count tokens in text."""
        if self._encoder:
            return len(self._encoder.encode(text))
        # Fallback: word-based approximation
        words = len(re.findall(r"\S+", text))
        return int(words / WORDS_PER_TOKEN)

    def analyze(self, content: str) -> dict:
        """
        Analyze token usage in CLAUDE.md.
        Returns breakdown by section and overall stats.
        """
        total = self.count(content)

        # Analyze by section
        sections = {}
        current_section = "preamble"
        current_content = []

        for line in content.split("\n"):
            if line.startswith("## ") or line.startswith("# "):
                if current_content:
                    sections[current_section] = self.count("\n".join(current_content))
                current_section = line.lstrip("#").strip()
                current_content = [line]
            else:
                current_content.append(line)

        if current_content:
            sections[current_section] = self.count("\n".join(current_content))

        # Sort sections by token count
        sections_sorted = dict(sorted(sections.items(), key=lambda x: -x[1]))

        return {
            "total": total,
            "by_section": sections_sorted,
            "status": self._get_status(total),
            "percentage_of_recommended": round((total / WARN_TOKENS) * 100, 1),
            "using_tiktoken": self._encoder is not None,
        }

    def _get_status(self, tokens: int) -> str:
        if tokens <= WARN_TOKENS:
            return "good"
        elif tokens <= CRITICAL_TOKENS:
            return "warning"
        else:
            return "critical"

    def format_report(self, content: str) -> str:
        """Format a human-readable token report."""
        data = self.analyze(content)
        lines = [
            f"Token count: {data['total']} / {WARN_TOKENS} recommended ({data['percentage_of_recommended']}%)",
            f"Status: {data['status'].upper()}",
        ]

        if data["status"] != "good":
            lines.append(f"\nTop sections by token usage:")
            for section, count in list(data["by_section"].items())[:5]:
                lines.append(f"  {section}: {count} tokens")
            lines.append("\nConsider trimming verbose sections or removing information the AI can infer from code.")

        if not data["using_tiktoken"]:
            lines.append("\n(Note: using word-based token approximation. Install tiktoken for accuracy.)")

        return "\n".join(lines)
