import json


def read_class_names(class_names_path: str):
  """Reads class names from text file.
  Supports .txt and .json.
  
  Args:
    class_names_path: `str`, path to json/txt file containing classes. 
      Text file should contain one class name per line.
      Json file should contain only one dictionary, `Mapping[int, str]`
  """
  
  names = {}
  if class_names_path.endswith('.txt'):
    with open(class_names_path, 'r') as data:
      for idx, name in enumerate(data):
        names[idx] = name.strip('\n')
  
  elif class_names_path.endswith('.json'):
    with open(class_names_path) as f:
      names = json.load(f)
    if type(list(names.keys())[0]) == str and type(list(names.values())[0]) == int:
      names = dict((v,k) for k,v in names.items())

  else:
    raise NotImplementedError('File type is not .txt or .json, path %s' %class_names_path)

  if type(list(names.keys())[0]) != int:
    raise ValueError('Loaded dict %s has wrong key type %s' %(
      class_names_path, type(list(names.keys())[0]))) 
  if type(list(names.values())[0]) != str:
    raise ValueError('Loaded dict %s has wrong value type %s' %(
      class_names_path, type(list(names.values())[0])))

  return names