# Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import boto3
import logging
import time
from datetime import datetime, timezone
from botocore.exceptions import WaiterError
from crhelper import CfnResource
from retrying import retry

logger = logging.getLogger(__name__)
cfnResource = CfnResource(json_logging=False, log_level='INFO', polling_interval=1)
iotsitewise = boto3.client('iotsitewise')

STANDARD_RETRY_MAX_ATTEMPT_COUNT = 10
WAITER_ERROR_RETRY_MAX_ATTEMPT_COUNT = 2
WAIT_TIME_IN_MILLISECONDS = 3000

@cfnResource.create
def create_asset_model(event, context):
    try:
        logger.info('Create asset model')
        logger.info(event)

        asset_model_id, metadata_sitewise_property_ids = create_asset_model_and_wait_until_active_with_retries(event)

        if metadata_sitewise_property_ids:
            update_property_ids_metadata_with_retries(asset_model_id, metadata_sitewise_property_ids)
        
        return asset_model_id
    except Exception as e:
        logger.exception(e)
        raise e

def is_waiter_error(exception):
    return isinstance(exception, WaiterError)

@retry(retry_on_exception=is_waiter_error,wait_fixed=WAIT_TIME_IN_MILLISECONDS,stop_max_attempt_number=WAITER_ERROR_RETRY_MAX_ATTEMPT_COUNT)
def create_asset_model_and_wait_until_active_with_retries(event):
    name = event['ResourceProperties']['name']
    description = event['ResourceProperties']['description']
    properties = event['ResourceProperties']['properties']
    hierarchies = event['ResourceProperties']['hierarchies']
    metadata_sitewise_property_ids = event['ResourceProperties']['metadata_sitewise_property_ids']
    composite_models = event['ResourceProperties']['assetModelCompositeModels']

    create_asset_model_response = create_asset_model_helper(name, description, properties, hierarchies, composite_models)
    asset_model_id = create_asset_model_response['assetModelId']
    iotsitewise.get_waiter(waiter_name='asset_model_active').wait(assetModelId=asset_model_id)
    
    return asset_model_id, metadata_sitewise_property_ids

def create_asset_model_helper(name, description, properties, hierarchies, composite_models):
    try:
        create_asset_model_response = iotsitewise.create_asset_model(assetModelName=name, assetModelDescription=description, assetModelProperties=properties, assetModelHierarchies=hierarchies, assetModelCompositeModels=composite_models)
    except Exception as e:
        create_asset_model_response = create_asset_model_uuid_name_with_retries(name,description,properties,hierarchies,composite_models)
    return create_asset_model_response

@retry(wait_fixed=WAIT_TIME_IN_MILLISECONDS, stop_max_attempt_number=STANDARD_RETRY_MAX_ATTEMPT_COUNT)
def create_asset_model_uuid_name_with_retries(name,description,properties,hierarchies,composite_models):
    return iotsitewise.create_asset_model(assetModelName=name + ' ' + get_timestamp(), assetModelDescription=description, assetModelProperties=properties, assetModelHierarchies=hierarchies, assetModelCompositeModels=composite_models)

@retry(wait_fixed=WAIT_TIME_IN_MILLISECONDS, stop_max_attempt_number=STANDARD_RETRY_MAX_ATTEMPT_COUNT)
def update_property_ids_metadata_with_retries(asset_model_id, metadata_sitewise_property_ids):
    desc_response = iotsitewise.describe_asset_model(assetModelId=asset_model_id)
    asset_model_properties = desc_response['assetModelProperties']
    property_ids = get_property_ids(metadata_sitewise_property_ids, asset_model_properties)
    cfnResource.Data.update(property_ids)

@cfnResource.update
def no_op(_, __):
    pass

@cfnResource.delete
def delete_asset_model(event, context):
    try:
        asset_model_id = event['PhysicalResourceId']
        time.sleep(60)
        iotsitewise.delete_asset_model(assetModelId=asset_model_id)
        iotsitewise.get_waiter(waiter_name='asset_model_not_exists').wait(assetModelId=asset_model_id)
    except Exception as e:
        logger.error(e)
    finally:
        return event.get('PhysicalResourceId', None)

def get_property_ids(metadata_sitewise_property_ids, asset_model_properties):
    try:
        property_ids = {}
        if not metadata_sitewise_property_ids or not asset_model_properties:
            return property_ids
        for asset_model_property in asset_model_properties:
            if asset_model_property['name'] in metadata_sitewise_property_ids:
                key = asset_model_property['name'] + " id"
                property_ids[key] = asset_model_property['id']
        return property_ids
    except Exception as e:
        logger.error(e)
        raise e

def get_timestamp():
    return datetime.now(timezone.utc).strftime("%Y-%b-%d %H:%M:%S")

def handler(event, context):    
    cfnResource(event, context)