from typing import OrderedDict
import six


def BuildPartialUpdate(clear, remove_keys, set_entries, field_mask_prefix,
                       entry_cls, env_builder):
  """Builds the field mask and patch environment for an environment update.

  Follows the environments update semantic which applies operations
  in an effective order of clear -> remove -> set.

  Leading and trailing whitespace is stripped from elements in remove_keys
  and the keys of set_entries.

  Args:
    clear: bool, If true, the patch removes existing keys.
    remove_keys: iterable(string), Iterable of keys to remove.
    set_entries: {string: string}, Dict containing entries to set.
    field_mask_prefix: string, The prefix defining the path to the base of the
      proto map to be patched.
    entry_cls: AdditionalProperty, The AdditionalProperty class for the type
      of entry being updated.
    env_builder: [AdditionalProperty] -> Environment, A function which produces
      a patch Environment with the given list of entry_cls properties.


  Returns:
    (string, Environment), a 2-tuple of the field mask defined by the arguments
    and a patch environment produced by env_builder.
  """
  remove_keys = set(k.strip() for k in remove_keys or [])
  # set_entries is sorted by key to make it easier for tests to set the
  # expected patch object.
  set_entries = OrderedDict(
      (k.strip(), v) for k, v in sorted(six.iteritems(set_entries or {})))
  if clear:
    entries = [
        entry_cls(key=key, value=value)
        for key, value in six.iteritems(set_entries)
    ]
    return field_mask_prefix, env_builder(entries)

  field_mask_entries = []
  seen_keys = set()
  for key in remove_keys:
    field_mask_entries.append('{}.{}'.format(field_mask_prefix, key))
    seen_keys.add(key)
  entries = []
  for key, value in six.iteritems(set_entries):
    entries.append(entry_cls(key=key, value=value))
    if key not in seen_keys:
      field_mask_entries.append('{}.{}'.format(field_mask_prefix, key))
  # Sorting field mask entries makes it easier for tests to set the expected
  # field mask since dictionary iteration order is undefined.
  field_mask_entries.sort()
  return ','.join(field_mask_entries), env_builder(entries)