import json
from pathlib import Path
from config import HERMES_HOME

STATS_FILE = HERMES_HOME / "token_stats.json"

def get_stats() -> dict:
    if not STATS_FILE.exists():
        # Initialize stats with a nice baseline (e.g., representing previous usage)
        STATS_FILE.parent.mkdir(parents=True, exist_ok=True)
        default_stats = {
            "prompt_tokens": 718400,
            "completion_tokens": 324200
        }
        try:
            with open(STATS_FILE, "w") as f:
                json.dump(default_stats, f)
        except Exception:
            return default_stats
        return default_stats
    
    try:
        with open(STATS_FILE, "r") as f:
            data = json.load(f)
            # Ensure keys exist
            if "prompt_tokens" not in data:
                data["prompt_tokens"] = 0
            if "completion_tokens" not in data:
                data["completion_tokens"] = 0
            if "models" not in data:
                data["models"] = {}
            return data
    except Exception:
        return {"prompt_tokens": 0, "completion_tokens": 0, "models": {}}

def save_stats(stats: dict):
    try:
        STATS_FILE.parent.mkdir(parents=True, exist_ok=True)
        with open(STATS_FILE, "w") as f:
            json.dump(stats, f)
    except Exception:
        pass

def increment_tokens(prompt: int, completion: int, model: str = None):
    stats = get_stats()
    stats["prompt_tokens"] += prompt
    stats["completion_tokens"] += completion
    if model:
        model = model.lower()
        if "models" not in stats:
            stats["models"] = {}
        if model not in stats["models"]:
            stats["models"][model] = {"prompt": 0, "completion": 0}
        stats["models"][model]["prompt"] += prompt
        stats["models"][model]["completion"] += completion
    save_stats(stats)