"""State persistence manager for LocalRun."""
import logging
import os
import pickle

logger = logging.getLogger("localrun.state")

class StateManager:
    def __init__(self, data_dir: str | None = None):
        # If data_dir is provided, it pins it. Otherwise, it fetches from config.
        self._data_dir = data_dir

    def get_state_file(self) -> str | None:
        """Returns the path to the state file, or None if data_dir is not set."""
        from localrun.config import get_config
        data_dir = self._data_dir or get_config().data_dir
        if not data_dir:
            return None
        os.makedirs(data_dir, exist_ok=True)
        return os.path.join(data_dir, "localrun_state.pkl")

    def save_state(self, engines: dict) -> bool:
        """Saves the current state of all engines to disk."""
        state_file = self.get_state_file()
        if not state_file:
            logger.warning("Cannot save state: LOCALRUN_DATA_DIR is not set")
            return False
            
        try:
            with open(state_file, "wb") as f:
                pickle.dump(engines, f, protocol=pickle.HIGHEST_PROTOCOL)
            logger.info("Successfully saved state to %s", state_file)
            return True
        except Exception as e:
            logger.error("Failed to save state: %s", e)
            return False

    def load_state(self, engines: dict) -> bool:
        """Loads state from disk and updates the engines dictionary in-place."""
        state_file = self.get_state_file()
        if not state_file:
            logger.warning("Cannot load state: LOCALRUN_DATA_DIR is not set")
            return False
            
        if not os.path.exists(state_file):
            logger.info("No state file found at %s", state_file)
            return False
            
        try:
            with open(state_file, "rb") as f:
                loaded_engines = pickle.load(f)
            
            # Update the existing engines dictionary in-place so references remain valid
            engines.clear()
            engines.update(loaded_engines)
            
            logger.info("Successfully loaded state from %s", state_file)
            return True
        except Exception as e:
            logger.error("Failed to load state: %s", e)
            return False
