"""
datazen - A child class for adding variable-data loading capabilities to the
          environment dataset.
"""

# built-in
from typing import List

# internal
from datazen import ROOT_NAMESPACE
from datazen.enums import DataType
from datazen.environment.base import BaseEnvironment
from datazen.load import LoadedFiles, DEFAULT_LOADS
from datazen.variables import load as load_variables


class VariableEnvironment(BaseEnvironment):
    """
    The variable-data loading environment mixin. Only requires the base
    environment capability to function.
    """

    def load_variables(
        self,
        var_loads: LoadedFiles = DEFAULT_LOADS,
        name: str = ROOT_NAMESPACE,
    ) -> dict:
        """Load variable data, resolve any un-loaded variable directories."""

        # determine directories that need to be loaded
        data_type = DataType.VARIABLE

        with self.lock:
            to_load = self.get_to_load(data_type, name)

            # load new data
            variable_data = self.get_data(data_type, name)
            if to_load:
                variable_data.update(load_variables(to_load, var_loads))
                self.update_load_state(data_type, to_load, name)

        return variable_data

    def add_variable_dirs(
        self,
        dir_paths: List[str],
        rel_path: str = ".",
        name: str = ROOT_NAMESPACE,
        allow_dup: bool = False,
    ) -> int:
        """
        Add variable-data directories, return the number of directories added.
        """

        return self.add_dirs(
            DataType.VARIABLE, dir_paths, rel_path, name, allow_dup
        )
