"""Best-effort JST daily counts. No disk/network I/O on the generation path."""

import json
import logging
import os
import threading
from datetime import datetime, timedelta, timezone
from pathlib import Path

LOGGER = logging.getLogger(__name__)
JST = timezone(timedelta(hours=9), "Asia/Tokyo")


class GenerationStats:
    """Single-process writer; deliberately not a distributed accounting DB."""

    def __init__(self, directory=None, interval=300):
        self.path = Path(directory) / "qwen-moss" / "daily-counts.json" if directory else None
        self.interval = interval
        self._counts = {}
        self._revision = self._saved_revision = 0
        self._loaded = False
        self._lock = threading.Lock()
        self._io_lock = threading.Lock()
        self._stop = threading.Event()
        self._thread = None

    def record(self, family, model, count=1, completed_at=None):
        if type(count) is not int or count <= 0:
            return
        if family not in ("qwen", "moss") or not isinstance(model, str) or not model or len(model) > 256:
            raise ValueError("invalid model identifier")
        day = (completed_at or datetime.now(JST)).astimezone(JST).date().isoformat()
        with self._lock:
            models = self._counts.setdefault(day, {}).setdefault(family, {})
            models[model] = models.get(model, 0) + count
            self._revision += 1

    def snapshot(self):
        with self._lock:
            return self._snapshot_unlocked()

    def _load(self):
        if self._loaded or self.path is None:
            return
        try:
            if self.path.stat().st_size > 10_000_000:
                raise ValueError("statistics file too large")
            raw = json.loads(self.path.read_text(encoding="utf-8"))
            if raw.get("version") != 1 or not isinstance(raw.get("counts"), dict):
                raise ValueError("invalid statistics schema")
            if raw.get("timezone") != "Asia/Tokyo":
                raise ValueError("invalid statistics timezone")
            restored = {}
            for day, families in raw["counts"].items():
                datetime.strptime(day, "%Y-%m-%d")
                restored[day] = {}
                for family, models in families.items():
                    if family not in ("qwen", "moss") or not isinstance(models, dict):
                        raise ValueError("invalid model family")
                    restored[day][family] = {}
                    for model, count in models.items():
                        if (not isinstance(model, str) or not model or len(model) > 256
                                or type(count) is not int or not 0 <= count <= 10**12):
                            raise ValueError("invalid model count")
                        restored[day][family][model] = count
        except FileNotFoundError:
            restored = {}
        with self._lock:
            # Counts generated during startup loading must not be lost.
            for day, families in restored.items():
                for family, models in families.items():
                    target = self._counts.setdefault(day, {}).setdefault(family, {})
                    for model, count in models.items():
                        target[model] = target.get(model, 0) + count
            self._loaded = True

    def flush(self):
        if self.path is None:
            return
        with self._io_lock:
            self._load()
            with self._lock:
                revision = self._revision
                if revision == self._saved_revision:
                    return
                data = self._snapshot_unlocked()
            # Never hold the counter lock while accessing mounted storage.
            self.path.parent.mkdir(parents=True, exist_ok=True)
            temporary = self.path.with_suffix(".json.tmp")
            temporary.write_text(json.dumps(data, ensure_ascii=False, separators=(",", ":")), encoding="utf-8")
            os.replace(temporary, self.path)
            with self._lock:
                self._saved_revision = revision

    def _snapshot_unlocked(self):
        return {"version": 1, "timezone": "Asia/Tokyo",
                "counts": {day: {family: dict(models) for family, models in families.items()}
                           for day, families in self._counts.items()}}

    def start(self):
        if self._thread is not None:
            return
        if self.path is None:
            LOGGER.warning("GENERATION_STATS_DIR is unset; counts are memory-only until a persistent directory is configured.")
            return

        def worker():
            try:
                with self._io_lock:
                    self._load()
            except Exception:
                LOGGER.warning("Statistics restore failed; will retry without blocking generation.", exc_info=True)
            while not self._stop.wait(self.interval):
                try:
                    self.flush()
                except Exception:
                    LOGGER.warning("Statistics save failed; will retry at the next interval.", exc_info=True)

        self._thread = threading.Thread(target=worker, name="generation-statistics", daemon=True)
        self._thread.start()

    def stop(self):
        self._stop.set()
        if self._thread is not None:
            self._thread.join(timeout=1)


stats = GenerationStats(os.environ.get("GENERATION_STATS_DIR"))