import kubernetes
from kubernetes.client.models import V1Volume


class KubernetesVolume:
    """KubernetesVolume data passing method involves passing data by mounting a
    single multi-write Kubernetes volume to containers instead of using Argo's
    artifact passing method (which stores the data in an S3 blob store)."""

    def __init__(self, volume: V1Volume, path_prefix: str = 'artifact_data/'):
        if not isinstance(volume, (dict, V1Volume)):
            raise TypeError('volume must be either V1Volume or dict')
        self._volume = volume
        self._path_prefix = path_prefix

    def transform_workflow(self, workflow: dict) -> dict:
        from ..compiler._data_passing_using_volume import rewrite_data_passing_to_use_volumes
        if isinstance(self._volume, dict):
            volume_dict = self._volume
        else:
            volume_dict = kubernetes.kubernetes.client.ApiClient(
            ).sanitize_for_serialization(self._volume)
        return rewrite_data_passing_to_use_volumes(workflow, volume_dict,
                                                   self._path_prefix)

    def __call__(self, workflow: dict) -> dict:
        return self.transform_workflow(workflow)