state.py

"""状态层: 单体全局状态 + 锁 + 原子落盘 + HF Dataset 持久化.

设计要点:
- threading.Lock 全局唯一;限流检查与递增在同一把锁内
- 写盘: 临时文件 + os.replace 原子改名;CommitScheduler 5min 周期
- 关键事件(蜕皮/生病/康复)调用 scheduler.trigger() 立即提交
- 启动时显式从 Dataset 拉最新状态(Space 重启不丢命)
- snapshots: 每次喂养一帧; events: 日记与回放的史料
"""
from __future__ import annotations

import json
import os
import threading
import time
from datetime import datetime, timedelta, timezone

from traits import TRAIT_KEYS

DATA_DIR = os.environ.get("CRICKET_DATA", "data")
STATE_PATH = os.path.join(DATA_DIR, "state.json")
DATASET_REPO = os.environ.get("CRICKET_DATASET", "")  # e.g. "lerp/cricket-commons-state"
DAILY_FEED_LIMIT = int(os.environ.get("CRICKET_DAILY_LIMIT", "10"))

CST = timezone(timedelta(hours=8))  # UTC+8 切日

LOCK = threading.Lock()
_scheduler = None


def _today() -> str:
    return datetime.now(CST).strftime("%Y-%m-%d")


def new_state() -> dict:
    now = time.time()
    return {
        "born_at": now,
        "name": "未命名小蛐蛐",
        "traits": {k: 0.5 for k in TRAIT_KEYS},
        "momentum": {k: 0.0 for k in TRAIT_KEYS},
        "mood": "calm",
        "feed_count": 0,
        "sick": False,
        "sick_since": None,
        "recover_progress": 0,
        "diary": [],      # [{date, zh, en}]
        "battles": [],    # [{ts, challenger_text, reply, win}]
        "snapshots": [],  # [{ts, traits, svg_params}]
        "events": [],     # [{ts, type, input_digest, reaction_digest}]
        "feeders": {},    # {fingerprint: {date, count}} 软限流
    }


def load_state() -> dict:
    """启动流程: 先尝试从 Dataset 拉最新,再读本地,最后新生."""
    os.makedirs(DATA_DIR, exist_ok=True)
    if DATASET_REPO and not os.path.exists(STATE_PATH):
        try:
            from huggingface_hub import hf_hub_download
            p = hf_hub_download(repo_id=DATASET_REPO, filename="state.json",
                                repo_type="dataset", local_dir=DATA_DIR)
            print(f"[state] restored from dataset: {p}")
        except Exception as e:
            print(f"[state] no remote state ({e}); fresh born")
    if os.path.exists(STATE_PATH):
        with open(STATE_PATH, encoding="utf-8") as f:
            return json.load(f)
    s = new_state()
    save_state(s)
    return s


def _scheduler_lock():
    global _scheduler
    if _scheduler is not None:
        return _scheduler.lock
    return threading.Lock()


def save_state(state: dict) -> None:
    """临时文件 + 原子改名;若 scheduler 存在则在其 lock 内写."""
    os.makedirs(DATA_DIR, exist_ok=True)
    tmp = STATE_PATH + ".tmp"
    lock = _scheduler.lock if _scheduler is not None else threading.Lock()
    with lock:
        with open(tmp, "w", encoding="utf-8") as f:
            json.dump(state, f, ensure_ascii=False)
        os.replace(tmp, STATE_PATH)


def start_scheduler() -> None:
    """CommitScheduler: 每 5 分钟把 data/ 推到 Dataset repo."""
    global _scheduler
    if not DATASET_REPO:
        print("[state] CRICKET_DATASET unset; persistence disabled (local only)")
        return
    from huggingface_hub import CommitScheduler
    _scheduler = CommitScheduler(
        repo_id=DATASET_REPO,
        repo_type="dataset",
        folder_path=DATA_DIR,
        every=5,  # minutes
    )
    print(f"[state] CommitScheduler started → {DATASET_REPO}")


def trigger_commit() -> None:
    """关键事件(蜕皮/生病/康复)立即提交."""
    if _scheduler is not None:
        try:
            _scheduler.trigger()
        except Exception as e:
            print(f"[state] trigger failed: {e}")


def check_and_count_feed(state: dict, fingerprint: str) -> bool:
    """软限流: 同一把锁内检查+递增(调用方需已持有 LOCK).
    返回 True=允许喂养."""
    today = _today()
    rec = state["feeders"].get(fingerprint)
    if rec is None or rec.get("date") != today:
        rec = {"date": today, "count": 0}
    if rec["count"] >= DAILY_FEED_LIMIT:
        return False
    rec["count"] += 1
    state["feeders"][fingerprint] = rec
    # 防 feeders 无限膨胀: 只留今天的
    if len(state["feeders"]) > 5000:
        state["feeders"] = {k: v for k, v in state["feeders"].items() if v.get("date") == today}
    return True


def append_event(state: dict, etype: str, input_digest: str = "", reaction_digest: str = "") -> None:
    state["events"].append({
        "ts": time.time(), "type": etype,
        "input": (input_digest or "")[:80], "reaction": (reaction_digest or "")[:80],
    })


def append_snapshot(state: dict, svg_params: dict) -> None:
    state["snapshots"].append({"ts": time.time(), **svg_params})


def age_days(state: dict) -> float:
    return (time.time() - state.get("born_at", time.time())) / 86400.0


def events_for_day(state: dict, date_str: str) -> list[dict]:
    """某天(UTC+8)的事件,日记的史料."""
    out = []
    for e in state.get("events", []):
        d = datetime.fromtimestamp(e["ts"], CST).strftime("%Y-%m-%d")
        if d == date_str:
            out.append(e)
    return out