import nirvana_dl


class IOWrapper():

    def __init__(self,
                 io_type: str,
                 input_dir: str,
                 checkpoint_dir: str,
                 logs_dir: str,
                 output_dir: str) -> None:
        super().__init__()

        self.io_type = io_type
        self.input_dir=input_dir
        self.checkpoint_dir=checkpoint_dir
        self.logs_dir=logs_dir
        self.output_dir=output_dir

    def get_input_dir(self) :
        return self.input_dir

    def get_checkpoint_dir(self) :
        return self.checkpoint_dir

    def get_logs_dir(self) :
        return self.logs_dir

    def get_output_dir(self) :
        return self.output_dir


def get_nirvana_wrapper():
    return IOWrapper("nirvana",
                     nirvana_dl.input_data_path(),
                     nirvana_dl.snapshot.get_snapshot_path(),
                     nirvana_dl.logs_path(),
                     nirvana_dl.output_data_path())


def get_local_wrapper():
    return IOWrapper("local",
                     "input",
                     "checkpoint",
                     "logs",
                     "output")
