src/soci/persistence/snapshots.py

"""Snapshots — save and load full simulation state."""

from __future__ import annotations

import json
import logging
import os
from pathlib import Path
from typing import Optional, TYPE_CHECKING

if TYPE_CHECKING:
    from soci.engine.simulation import Simulation
    from soci.engine.llm import ClaudeClient
    from soci.persistence.database import Database

logger = logging.getLogger(__name__)

SNAPSHOTS_DIR = Path(os.environ.get("SOCI_DATA_DIR", "data")) / "snapshots"


async def save_simulation(
    sim: Simulation,
    db: Database,
    name: str = "autosave",
) -> None:
    """Save the full simulation state to database and JSON file."""
    state = sim.to_dict()

    # Save to database
    await db.save_snapshot(
        name=name,
        tick=sim.clock.total_ticks,
        day=sim.clock.day,
        state=state,
    )

    # Also save as JSON file for easy inspection
    SNAPSHOTS_DIR.mkdir(parents=True, exist_ok=True)
    path = SNAPSHOTS_DIR / f"{name}.json"
    with open(path, "w", encoding="utf-8") as f:
        json.dump(state, f, indent=2, ensure_ascii=False)

    logger.info(f"Simulation saved: {name} (tick {sim.clock.total_ticks}, day {sim.clock.day})")


async def load_simulation(
    db: Database,
    llm: ClaudeClient,
    name: Optional[str] = None,
) -> Optional[Simulation]:
    """Load a simulation from the database."""
    from soci.engine.simulation import Simulation

    state = await db.load_snapshot(name)
    if not state:
        # Fall back to JSON file (autosave.json when no name given)
        json_name = name or "autosave"
        path = SNAPSHOTS_DIR / f"{json_name}.json"
        if path.exists():
            try:
                with open(path, "r", encoding="utf-8") as f:
                    state = json.load(f)
                logger.info(f"Loaded snapshot from JSON fallback: {path}")
            except (json.JSONDecodeError, ValueError) as e:
                logger.warning(f"Snapshot file corrupt or empty, ignoring: {path} ({e})")
                state = None

    if not state:
        logger.info(f"No snapshot found: {name or 'latest'} — will start fresh")
        return None

    sim = Simulation.from_dict(state, llm)
    logger.info(f"Simulation loaded: tick {sim.clock.total_ticks}, day {sim.clock.day}")
    return sim