"""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")


def validate_counts(counts):
    if not isinstance(counts, dict):
        raise ValueError("invalid statistics counts")
    for day, families in counts.items():
        if datetime.strptime(day, "%Y-%m-%d").date().isoformat() != day or not isinstance(families, dict):
            raise ValueError("invalid statistics day")
        for family, models in families.items():
            if family not in ("anima", "illustrious") or not isinstance(models, dict):
                raise ValueError("invalid model 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")
    return counts


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

    def __init__(self, directory=None, interval=300):
        self.path = Path(directory) / "daily-counts.json" if directory else None
        self.interval = interval
        self._counts = {}
        self._api_counts = {}
        self._api_started_at = datetime.now(JST).isoformat()
        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, source="webui"):
        if type(count) is not int or count <= 0:
            return
        if source not in ("webui", "api"):
            raise ValueError("invalid generation source")
        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
            if source == "api":
                api_models = self._api_counts.setdefault(day, {}).setdefault(family, {})
                api_models[model] = api_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")
            restored = validate_counts(raw["counts"])
            api_restored = validate_counts(raw.get("api_counts", {}))
            started_at = raw.get("api_tracking_started_at")
            if started_at is not None:
                stamp = datetime.fromisoformat(started_at)
                if stamp.tzinfo is None:
                    raise ValueError("invalid API tracking start")
                started_at = stamp.astimezone(JST).isoformat()
            elif api_restored:
                raise ValueError("API counts without tracking start")
            for day, families in api_restored.items():
                for family, models in families.items():
                    for model, count in models.items():
                        if count > restored.get(day, {}).get(family, {}).get(model, 0):
                            raise ValueError("API count exceeds total")
        except FileNotFoundError:
            restored = {}
            api_restored = {}
            started_at = None
        with self._lock:
            # Counts generated during startup loading must not be lost.
            for saved, current in ((restored, self._counts), (api_restored, self._api_counts)):
                for day, families in saved.items():
                    for family, models in families.items():
                        target = current.setdefault(day, {}).setdefault(family, {})
                        for model, count in models.items():
                            target[model] = target.get(model, 0) + count
            if started_at:
                self._api_started_at = min(self._api_started_at, started_at)
            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()},
                "api_tracking_started_at": self._api_started_at,
                "api_counts": {day: {family: dict(models) for family, models in families.items()}
                               for day, families in self._api_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"))