from typing import Dict
from typing import Any


def _get_nested_metadata(dataset: Dict[str, Any],
                         prefix: str) -> Dict[str, Any]:
  """Generate a metadata dictionary using flattened metadata keys.

  Args:
    dataset: dictionary containing the dataset keys and values. Keys are
      flatened.
    prefix: common prefix of the metadata fields.

  Returns:
    Nested dictionary with the episode metadata.

  If the dataset contains:
  {
    'metadata/v1/v2': 1,
    'metadata/v3': 2,
  }
  and prefix='metadata', it returns:
  {
    'v1':{
      'v2': 1,
    }
    'v3': 2,
  }
  It assumes that the flattened metadata keys are well-formed.
  """
  episode_metadata = {}
  for k in dataset.keys():
    if f'{prefix}/' not in k:
      continue
    keys = k.split('/')[1:]
    nested_dict = episode_metadata
    leaf_value = dataset[k]
    for index, nested_key in enumerate(keys):
      if index == (len(keys) - 1):
        nested_dict[nested_key] = leaf_value
      else:
        if nested_key not in nested_dict:
          nested_dict[nested_key] = {}
        nested_dict = nested_dict[nested_key]

  return episode_metadata