Datastore

@dataclass
class DataStore:
    """A data store for managing optimization-related data with multiple storage types.

    This class provides a unified interface for storing and managing different types of
    optimization-related data, including problems, instances, solutions, and various
    metadata. It supports both file system storage and OMMX artifact storage.

    The data store maintains separate storage for:
    - Problems: JijModeling Problem instances
    - Instances: OMMX Instance objects
    - Solutions: OMMX Solution objects
    - Objects: Generic JSON-serializable objects
    - Parameters: Configuration parameters
    - Samplesets: OMMX SampleSet objects
    - Meta-data: Additional metadata

    Directory structure:
    ```
    dir
     ├── problem_*.problem    # Individual problem files
     ├── instance_*.instance  # Individual instance files
     ├── solution_*.solution  # Individual solution files
     ├── objects_*.json      # Individual object files
     ├── parameters_.json    # Single parameters file
     ├── samplesets_*.sampleset  # Individual sampleset files
     └── meta_data_.json     # Single metadata file
    ```

    Attributes:
        problems (dict[str, jm.Problem]): Storage for optimization problems
        instances (dict[str, ommx_v1.Instance]): Storage for problem instances
        solutions (dict[str, ommx_v1.Solution]): Storage for problem solutions
        objects (dict[str, dict]): Storage for generic JSON-serializable objects
        parameters (dict[str, dict[str, Any]]): Storage for parameters
        meta_data (dict[str, Any]): Storage for metadata
    """

    problems: dict[str, jm.Problem] = field(default_factory=dict)
    instances: dict[str, ommx_v1.Instance] = field(default_factory=dict)
    solutions: dict[str, ommx_v1.Solution] = field(default_factory=dict)
    objects: dict[str, dict] = field(default_factory=dict)
    parameters: dict[str, float | int] = field(default_factory=dict)
    samplesets: dict[str, ommx_v1.SampleSet] = field(default_factory=dict)
    meta_data: dict[str, typ.Any] = field(default_factory=dict)

    _storage_mapping: typ.ClassVar[dict[str, StorageStrategy]] = {
        "problems": ProblemStorage(),
        "instances": InstanceStorage(),
        "solutions": SolutionStorage(),
        "objects": JSONStorage(),
        "parameters": JSONStorage(),
        "samplesets": SampleSetStorage(),
        "meta_data": JSONStorage(),
    }

    minto_namespace: typ.ClassVar[str] = "org.minto"

    def add(
        self,
        name: str,
        obj,
        storage_name: str,
        with_save=False,
        save_dir: str | pathlib.Path = ".",
    ):
        """Add an object to the specified storage.

        Args:
            name (str): Identifier for the object
            obj (Any): Object to store
            storage_name (str): Type of storage ('problems', 'instances', etc.)
            with_save (bool, optional): Whether to save to disk. Defaults to False.
            save_dir (str, optional): Directory for saving files. Defaults to ".".

        Examples:
            >>> ds = DataStore()
            >>> ds.add("problem1", problem, "problems", with_save=True)
            >>> ds.add("param1", {"value": 42}, "parameters")
        """
        # Add the object to the storage
        getattr(self, storage_name)[name] = obj

        # Save the object to the file
        if with_save:
            # check existance of save_dir
            save_dir = pathlib.Path(save_dir)
            if not save_dir.exists():
                save_dir.mkdir(parents=True)
            match storage_name:
                case "parameters" | "meta_data":
                    # Save the whole dictionary to the file
                    # Example:
                    # (original) .parameters = {"a": 1, "b": 2}
                    # user input: add("new", 3, "parameters")
                    # (new) .parameters = {"a": 1, "b": 2, "new": 3}
                    # save to the file: {"a": 1, "b": 2, "new": 3}
                    storage = self._storage_mapping[storage_name]
                    file_name = self._file_name(
                        storage_name, "", storage.extension
                    )  # noqa
                    attr_value = getattr(self, storage_name)
                    storage.save(attr_value, save_dir / pathlib.Path(file_name))  # noqa
                case _:
                    # save each the object to the file with the name
                    # Example:
                    # (original) .objects = {"a": {"x": 1}}
                    # user input: add("new", {"z": 3}, "objects")
                    # (new) .objects = {"a": {"x": 1}, "new": {"z": 3}}
                    # save to the file (objects_new.json): {"z": 3}
                    storage = self._storage_mapping[storage_name]
                    file_name = self._file_name(
                        storage_name, name, storage.extension
                    )  # noqa
                    storage.save(obj, save_dir / pathlib.Path(file_name))

    @classmethod
    def _file_name(cls, name: str, obj_name: str, extionsion: str):
        return f"{name}_{obj_name}.{extionsion}"

    def save_all(self, path: pathlib.Path):
        """Save all stored data to the specified directory.

        Saves all objects in all storage types to their respective files in the
        specified directory, maintaining the standard directory structure.

        Args:
            path (pathlib.Path): Directory where files should be saved
        """

        # Check path is exist or not
        if not path.exists():
            path.mkdir(parents=True)
        for storage_name, storage in self._storage_mapping.items():
            storage: StorageStrategy
            attr_value: dict = getattr(self, storage_name)
            if storage_name in ("parameters", "meta_data"):
                # Save the whole dictionary to the file
                file_name = self._file_name(storage_name, "", storage.extension)
                storage.save(attr_value, path / file_name)
            else:
                # save each the object to the file with the name
                for obj_name, obj in attr_value.items():
                    file_name = self._file_name(
                        storage_name, obj_name, storage.extension
                    )
                    storage.save(obj, path / file_name)

    @classmethod
    def load(cls, path: pathlib.Path):
        """Load a DataStore from a directory.

        Creates a new DataStore instance and populates it with data from files
        in the specified directory.

        Args:
            path (pathlib.Path): Directory containing the stored files

        Returns:
            DataStore: New DataStore instance containing the loaded data
        """
        data = {}
        for storage_name, storage in cls._storage_mapping.items():
            storage: StorageStrategy
            if storage_name in ("parameters", "meta_data"):
                file_path = path / cls._file_name(
                    storage_name, "", storage.extension
                )  # noqa
                if file_path.exists():
                    data[storage_name] = storage.load(file_path)
            else:
                data[storage_name] = {}
                file_pattern = cls._file_name(storage_name, "*", storage.extension)
                for file_path in path.glob(str(file_pattern)):  # noqa
                    obj_name = "_".join(
                        file_path.stem.split(".")[0].split("_")[1:]
                    )  # noqa
                    data[storage_name][obj_name] = storage.load(file_path)
        return cls(**data)

    def add_to_artifact_builder(
        self, builder: ox_art.ArtifactBuilder, annotations: dict
    ):
        """Add all stored data to an OMMX artifact builder.

        Adds all objects from all storage types to the artifact builder, including
        appropriate annotations for each layer.

        Args:
            builder (ox_art.ArtifactBuilder): The artifact builder
            annotations (dict[str, str]): Base annotations for all layers
        """
        _annotations = copy.copy(annotations)
        for storage_name, storage in self._storage_mapping.items():
            storage: StorageStrategy
            _annotations[self.minto_namespace + ".storage"] = storage_name
            if storage_name in ("parameters", "meta_data"):
                obj = getattr(self, storage_name)
                storage.add_to_artifact_builder(obj, builder, _annotations)
            else:
                for obj_name, obj in getattr(self, storage_name).items():
                    _annotations[self.minto_namespace + ".name"] = obj_name
                    storage.add_to_artifact_builder(obj, builder, _annotations)

    @classmethod
    def load_from_layers(
        cls, artifact: ox_art.Artifact, layers: typ.Iterable[ox_art.Descriptor]
    ):
        """Create a DataStore from OMMX artifact layers.

        Args:
            artifact (ox_art.Artifact): The OMMX artifact containing the data
            layers (Iterable[ox_art.Descriptor]): Layer descriptors to process

        Returns:
            DataStore: New DataStore instance containing the data from the layers
        """
        datastore = cls()
        for layer in layers:
            storage_name = layer.annotations.get(
                cls.minto_namespace + ".storage"
            )  # noqa
            if storage_name is None:
                continue
            storage = cls._storage_mapping[storage_name]
            obj = storage.load_from_layer(artifact, layer)
            if storage_name == "parameters":
                datastore.parameters = obj
            elif storage_name == "meta_data":
                datastore.meta_data = obj
            else:
                obj_name = layer.annotations.get(cls.minto_namespace + ".name")
                if obj_name is None:
                    continue
                getattr(datastore, storage_name)[obj_name] = obj
        return datastore

    def get_solution(self) -> ommx_v1.Solution:
        """Get a solution from runs.

        If there are multiple solutions, return the first solution sorted by key.
        This is a syntax sugar for the common use case where there is only one solution.
        If you have multiple solutions in one experiment or one run, refer to the .solutions property instead.

        Returns:
            ommx.v1.Solution: The first solution sorted by key
        """
        # get min key value of self.solutions
        sorted_keys: list[str] = sorted(self.solutions.keys())
        return self.solutions[sorted_keys[0]]

class StorageStrategy(abc.ABC, typ.Generic[T]):
    """Abstract base class defining the interface for storage strategies.

    This class provides a generic interface for storing and loading different types of data,
    supporting both file system storage and OMMX artifact storage.

    Type Parameters:
        T: The type of data this strategy handles
    """

    @abc.abstractmethod
    def save(self, data: T, path: pathlib.Path):
        """Save data to the specified path.

        Args:
            data (T): The data to save
            path (pathlib.Path): The path where the data should be saved
        """
        pass

    @abc.abstractmethod
    def load(self, path: pathlib.Path) -> T:
        """Load data from the specified path.

        Args:
            path (pathlib.Path): The path from which to load the data

        Returns:
            T: The loaded data
        """
        pass

    @abc.abstractmethod
    def add_to_artifact_builder(
        self, data, builder: ox_art.ArtifactBuilder, annotations: dict[str, str]
    ) -> None:
        pass

    @abc.abstractmethod
    def load_from_layer(self, artifact: ox_art.Artifact, layer: ox_art.Descriptor) -> T:
        pass

    @property
    @abc.abstractmethod
    def extension(self) -> str:
        pass