#!/usr/bin/env python3
"""
Token Estimator -- Estimate token costs for files and directories.

Helps agents and developers understand which files are expensive to read
and where token budgets are being spent. Uses a simple heuristic:
~0.25 tokens per character for code, ~0.20 for prose.

Usage:
    token-estimator src/                    # estimate tokens for directory
    token-estimator src/App.tsx             # estimate for single file
    token-estimator . --top 20             # top 20 most expensive files
    token-estimator . --budget 50000       # files that fit in 50K token budget
    token-estimator . --category            # group by file category
    token-estimator --version               # show version
"""

import sys
import os
import json
from pathlib import Path
from collections import defaultdict

__version__ = "1.0.0"

# ANSI colors
class C:
    GREEN = "\033[92m"
    YELLOW = "\033[93m"
    RED = "\033[91m"
    BLUE = "\033[94m"
    BOLD = "\033[1m"
    DIM = "\033[2m"
    RESET = "\033[0m"

    @staticmethod
    def disable():
        C.GREEN = C.YELLOW = C.RED = C.BLUE = C.BOLD = C.DIM = C.RESET = ""


# Token estimation ratios (tokens per character)
RATIOS = {
    'code': 0.25,       # TypeScript, JavaScript, Python, Rust
    'prose': 0.20,      # Markdown, plain text
    'data': 0.30,       # JSON, YAML (lots of punctuation = more tokens)
    'config': 0.28,     # Config files
}

EXTENSION_MAP = {
    '.ts': 'code', '.tsx': 'code', '.js': 'code', '.jsx': 'code',
    '.py': 'code', '.rs': 'code', '.go': 'code', '.rb': 'code',
    '.java': 'code', '.cpp': 'code', '.c': 'code', '.h': 'code',
    '.css': 'code', '.scss': 'code', '.less': 'code',
    '.md': 'prose', '.txt': 'prose', '.rst': 'prose',
    '.json': 'data', '.yaml': 'data', '.yml': 'data', '.toml': 'data',
    '.xml': 'data', '.html': 'data', '.svg': 'data',
    '.sh': 'config', '.bash': 'config', '.zsh': 'config',
    '.env': 'config', '.ini': 'config', '.cfg': 'config',
    '.gitignore': 'config', '.eslintrc': 'config',
}

SKIP_DIRS = {
    'node_modules', '.git', 'dist', 'build', '.next', '__pycache__',
    'target', '.cache', '.vault-search', '_github_imports',
}

CATEGORY_MAP = {
    '.claude/agents': 'Agent Definitions',
    '.claude/skills': 'Skills',
    '.claude/hooks': 'Hooks',
    'agent-learnings': 'Learning System',
    'handoffs': 'Handoffs',
    'docs/specs': 'Specifications',
    'docs/plans': 'Plans',
    'docs/guides': 'Guides',
    'docs/reference': 'Reference',
    'docs/research': 'Research',
    'databases': 'Databases',
    'scrapers': 'Scrapers',
    'utilities': 'Utilities',
    'components': 'Components',
    'stores': 'Stores',
    'lib': 'Libraries',
    'app': 'App Pages',
    'types': 'Types',
}


def estimate_tokens(filepath):
    """Estimate token count for a file."""
    try:
        content = filepath.read_text(encoding='utf-8', errors='ignore')
    except Exception:
        return 0

    chars = len(content)
    file_type = EXTENSION_MAP.get(filepath.suffix, 'code')
    ratio = RATIOS.get(file_type, 0.25)
    return int(chars * ratio)


def categorize_file(path_str):
    """Categorize a file by its directory."""
    for prefix, category in CATEGORY_MAP.items():
        if prefix in path_str:
            return category
    return 'Other'


def scan_directory(root, max_files=None):
    """Scan directory and estimate tokens for all files."""
    root = Path(root).resolve()
    results = []

    for filepath in sorted(root.rglob('*')):
        if any(skip in filepath.parts for skip in SKIP_DIRS):
            continue
        if not filepath.is_file():
            continue
        if filepath.suffix not in EXTENSION_MAP:
            continue
        if filepath.stat().st_size > 5_000_000:  # Skip files >5MB
            continue

        tokens = estimate_tokens(filepath)
        if tokens == 0:
            continue

        rel_path = str(filepath.relative_to(root))
        results.append({
            'path': rel_path,
            'tokens': tokens,
            'chars': filepath.stat().st_size,
            'lines': filepath.read_text(encoding='utf-8', errors='ignore').count('\n') + 1,
            'type': EXTENSION_MAP.get(filepath.suffix, 'code'),
            'category': categorize_file(rel_path),
        })

    results.sort(key=lambda x: -x['tokens'])
    if max_files:
        results = results[:max_files]

    return results


def format_tokens(n):
    """Format token count with K suffix."""
    if n >= 1000:
        return f"{n/1000:.1f}K"
    return str(n)


def print_results(results, budget=None, by_category=False):
    """Print token estimates."""
    if not results:
        print("No files found.")
        return

    total_tokens = sum(r['tokens'] for r in results)

    print(f"\n{C.BOLD}Token Estimates{C.RESET}")
    print(f"{C.DIM}{'=' * 60}{C.RESET}\n")

    if by_category:
        categories = defaultdict(lambda: {'files': 0, 'tokens': 0})
        for r in results:
            cat = r['category']
            categories[cat]['files'] += 1
            categories[cat]['tokens'] += r['tokens']

        print(f"  {'Category':<25s} {'Files':>6s} {'Tokens':>10s} {'%':>6s}")
        print(f"  {'-' * 25} {'-' * 6} {'-' * 10} {'-' * 6}")

        for cat, stats in sorted(categories.items(), key=lambda x: -x[1]['tokens']):
            pct = stats['tokens'] / total_tokens * 100 if total_tokens > 0 else 0
            color = C.RED if pct > 30 else C.YELLOW if pct > 15 else ""
            print(f"  {cat:<25s} {stats['files']:>6d} {color}{format_tokens(stats['tokens']):>10s}{C.RESET} {pct:>5.1f}%")

    else:
        if budget:
            print(f"  {C.BLUE}Budget: {format_tokens(budget)} tokens{C.RESET}")
            print()

        running_total = 0
        budget_exceeded = False

        for r in results:
            running_total += r['tokens']
            if budget and running_total > budget and not budget_exceeded:
                print(f"  {C.RED}--- budget exceeded at {format_tokens(running_total)} ---{C.RESET}")
                budget_exceeded = True

            pct = r['tokens'] / total_tokens * 100 if total_tokens > 0 else 0
            bar_len = int(pct / 2)
            bar = '#' * bar_len

            color = C.RED if r['tokens'] > 5000 else C.YELLOW if r['tokens'] > 2000 else ""
            dim = C.DIM if budget_exceeded else ""

            print(f"  {dim}{color}{format_tokens(r['tokens']):>7s}{C.RESET}{dim}  {bar:<25s}  {r['path']}{C.RESET}")

    print(f"\n{C.DIM}{'=' * 60}{C.RESET}")
    print(f"  Total: {C.BOLD}{format_tokens(total_tokens)}{C.RESET} tokens across {len(results)} files")

    if budget:
        remaining = budget - total_tokens
        if remaining > 0:
            print(f"  Budget remaining: {C.GREEN}{format_tokens(remaining)}{C.RESET}")
        else:
            print(f"  Over budget by: {C.RED}{format_tokens(-remaining)}{C.RESET}")

    print()


def main():
    args = sys.argv[1:]

    if "--version" in args:
        print(f"token-estimator {__version__}")
        sys.exit(0)

    if "--help" in args or not args:
        print(__doc__)
        sys.exit(0)

    if not sys.stdout.isatty():
        C.disable()

    # Parse flags
    top_n = None
    budget = None
    by_category = "--category" in args
    as_json = "--json" in args

    if "--top" in args:
        idx = args.index("--top")
        if idx + 1 < len(args):
            top_n = int(args[idx + 1])

    if "--budget" in args:
        idx = args.index("--budget")
        if idx + 1 < len(args):
            budget = int(args[idx + 1])

    # Get target
    target = None
    for a in args:
        if not a.startswith("--") and a not in (str(top_n), str(budget)):
            target = a
            break

    if not target:
        target = "."

    target_path = Path(target).resolve()

    if target_path.is_file():
        tokens = estimate_tokens(target_path)
        chars = target_path.stat().st_size
        print(f"\n  {target}: ~{format_tokens(tokens)} tokens ({chars:,} chars)")
        print()
    elif target_path.is_dir():
        results = scan_directory(target_path, max_files=top_n)
        if as_json:
            print(json.dumps(results, indent=2))
        else:
            print_results(results, budget, by_category)
    else:
        print(f"Not found: {target}")
        sys.exit(1)


if __name__ == "__main__":
    main()