#!/usr/bin/env python3
"""
token_optimizer.py — Token Efficiency & Optimization for the Timmy Foundation.

Analyzes agent logs to identify:
1. "Chatty" Agents — agents outputting excessive tokens for simple tasks.
2. Redundant Logs — identifying patterns of repetitive log output.
3. Tool Output Bloat — identifying tools that return unnecessarily large payloads.

Outputs an "Efficiency Score" (0-100) per agent.
"""

import os
import sys
import glob
import re
from pathlib import Path
from collections import defaultdict
from typing import Dict, List

AGENT_LOG_PATHS = [
    "/root/wizards/*/home/logs/*.log",
    "/root/wizards/*/logs/*.log",
    "/root/wizards/*/.hermes/logs/*.log",
]

class TokenOptimizer:
    def __init__(self):
        self.agent_stats = defaultdict(lambda: {"tokens": 0, "turns": 0, "tool_calls": 0})

    def estimate_tokens(self, text: str) -> int:
        # Rough estimate: 4 chars per token
        return len(text) // 4

    def find_logs(self) -> List[Path]:
        files = []
        for pattern in AGENT_LOG_PATHS:
            for p in glob.glob(pattern):
                files.append(Path(p))
        return files

    def analyze_log(self, path: Path):
        # Extract agent name from path
        try:
            parts = path.parts
            idx = parts.index("wizards")
            agent = parts[idx + 1]
        except (ValueError, IndexError):
            agent = "unknown"

        try:
            with open(path, "r", errors="ignore") as f:
                content = f.read()
                self.agent_stats[agent]["tokens"] += self.estimate_tokens(content)
                
                # Count turns (approximate by looking for role markers)
                self.agent_stats[agent]["turns"] += content.count("[ASSISTANT]")
                self.agent_stats[agent]["turns"] += content.count("[USER]")
                
                # Count tool calls
                self.agent_stats[agent]["tool_calls"] += content.count("Calling tool:")
        except Exception as e:
            print(f"Error analyzing {path}: {e}")

    def run(self):
        print("--- Token Efficiency Audit ---")
        logs = self.find_logs()
        for log in logs:
            self.analyze_log(log)

        print(f"{'Agent':<20} | {'Tokens':<10} | {'Turns':<6} | {'T/Turn':<8} | {'Efficiency'}")
        print("-" * 65)
        
        for agent, stats in self.agent_stats.items():
            tokens = stats["tokens"]
            turns = max(stats["turns"], 1)
            t_per_turn = tokens // turns
            
            # Efficiency score: lower tokens per turn is generally better
            # Baseline: 500 tokens per turn = 100 score. 2000+ = 0 score.
            efficiency = max(0, min(100, 100 - (t_per_turn - 500) // 15))
            
            print(f"{agent:<20} | {tokens:<10} | {turns:<6} | {t_per_turn:<8} | {efficiency}%")

if __name__ == "__main__":
    optimizer = TokenOptimizer()
    optimizer.run()