"""Token-Statistik (Verbrauch je Modell) mit gedrosseltem Persistieren.

Früher wurde bei JEDEM Request die komplette JSON-Datei gelesen und geschrieben
(Disk-Thrash). Jetzt: einmaliges Laden in einen In-Memory-Cache, Inkremente laufen
gegen den Cache, Persistieren passiert höchstens alle FLUSH_INTERVAL Sekunden sowie
beim Prozess-Ende (atexit). Lesen liefert immer den aktuellen (auch ungeflushten) Stand.
"""

import atexit
import json
import logging
import threading
import time
from pathlib import Path

from config import HERMES_HOME

STATS_FILE = HERMES_HOME / "token_stats.json"
FLUSH_INTERVAL = 5.0  # Sekunden zwischen Disk-Writes
# Baseline (repräsentiert Verbrauch vor dem modellspezifischen Logging).
_BASELINE = {"prompt_tokens": 718400, "completion_tokens": 324200, "models": {}}

log = logging.getLogger(__name__)

_lock = threading.Lock()
_stats: dict | None = None
_dirty = False
_last_flush = 0.0


def _load_from_disk() -> dict:
    if not STATS_FILE.exists():
        return dict(_BASELINE)
    try:
        with open(STATS_FILE, "r", encoding="utf-8") as f:
            data = json.load(f)
        data.setdefault("prompt_tokens", 0)
        data.setdefault("completion_tokens", 0)
        data.setdefault("models", {})
        return data
    except (OSError, json.JSONDecodeError):
        log.warning("token_stats: Laden fehlgeschlagen, nutze Baseline", exc_info=True)
        return dict(_BASELINE)


def _ensure_loaded() -> dict:
    global _stats
    if _stats is None:
        _stats = _load_from_disk()
    return _stats


def _write(stats: dict) -> None:
    try:
        STATS_FILE.parent.mkdir(parents=True, exist_ok=True)
        tmp = STATS_FILE.with_suffix(".tmp")
        with open(tmp, "w", encoding="utf-8") as f:
            json.dump(stats, f)
        tmp.replace(STATS_FILE)
    except OSError:
        log.warning("token_stats: Schreiben fehlgeschlagen", exc_info=True)


def get_stats() -> dict:
    """Aktueller Stand (inkl. noch nicht geflushter Inkremente) als Kopie."""
    with _lock:
        return json.loads(json.dumps(_ensure_loaded()))


def increment_tokens(prompt: int, completion: int, model: str | None = None) -> None:
    """Tokens im Cache verbuchen; gedrosselt auf Disk persistieren."""
    global _dirty, _last_flush
    with _lock:
        stats = _ensure_loaded()
        stats["prompt_tokens"] += prompt
        stats["completion_tokens"] += completion
        if model:
            m = stats.setdefault("models", {}).setdefault(
                model.lower(), {"prompt": 0, "completion": 0})
            m["prompt"] += prompt
            m["completion"] += completion
        _dirty = True
        now = time.monotonic()
        if now - _last_flush >= FLUSH_INTERVAL:
            _write(stats)
            _dirty = False
            _last_flush = now


def flush() -> None:
    """Ungeschriebene Inkremente sofort persistieren (z.B. beim Shutdown)."""
    global _dirty
    with _lock:
        if _dirty and _stats is not None:
            _write(_stats)
            _dirty = False


atexit.register(flush)