import boto3
from os import environ

'''
Configure the script from enviroment variables.
'''
riv_stack_name = environ.get('RIV_STACK_NAME')
region_name = environ.get('REGION')
total_collections = environ.get('TOTAL_COLLECTIONS')

if riv_stack_name is None:
  print('RIV_STACK_NAME variable missing, defaulting to Riv-Prod')
  riv_stack_name='Riv-Prod'

if total_collections is None:
  print('TOTAL_COLLECTIONS variable missing, defaulting to 10')
  total_collections=10

if region_name is None:
  for alter in ['AWS_REGION', 'AWS_DEFAULT_REGION', 'AMAZON_REGION']:
    region_name = environ.get(alter)
    if region_name is not None:
      print('Defaulting region_name to env[{%s}] = %s' % (alter, region_name))
      break

assert riv_stack_name is not None, "riv_stack_name is not available"
assert region_name is not None, "region_name is not available"
assert total_collections is not None, "region_name is not available"

rek_client = boto3.client('rekognition', region_name=region_name)
ssm_client = boto3.client('ssm', region_name=region_name)
total_collections = int(total_collections)

def create_collections()->None:
  '''
  Create each of the Rekognition Collection Partitions
  '''
  for ix in range(0, total_collections):
    collection_id = '{}-{}'.format(riv_stack_name, ix)

    try:
      '''
      Check if the collection exists, or skip forward if already present.
      '''
      rek_client.describe_collection(CollectionId=collection_id)
      print('Collection {} already exists, nothing to do.'.format(collection_id))
      continue
    except rek_client.exceptions.ResourceNotFoundException as error:
      '''
      This error means the collection does not exist. Create it.
      '''
      response = rek_client.create_collection(
        CollectionId=collection_id,
        Tags={
          'riv_stack': riv_stack_name,
        })

      print('Created collection {} with Version {}'.format(collection_id, response['FaceModelVersion']))
    except Exception as error:
      '''
      These are hard failures (e.g., AccessDeniedException)
      '''
      print('Unable to describe_collection({})'.format(collection_id))
      raise error

def set_parameters()->None:
  '''
  Set the parameters for the Rekogonion clients.
  '''
  for name, value in [('partition-count',str(total_collections))]:
    try:
      response = ssm_client.put_parameter(
        Name='/riv/{}/rekognition/{}'.format(riv_stack_name, name),
        Value=value,
        Description='Generated by {}'.format(__file__),
        Type='String',
        Overwrite=True,
        Tier='Standard',
        DataType='text')
      print('set_parameter({}) with version {}'.format(name, response['Version']))
    except Exception as error:
      print('Unable to set_parameter({}, {})'.format(name,value))
      raise error

if __name__ == '__main__':
  create_collections()
  set_parameters()
  print('Setup is complete.')