# 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(event, context):
    try:
        logger.info('Create asset')
        logger.info(event)

        asset_id, asset_ids_to_associate = create_asset_and_wait_until_active_with_retries(event)

        if asset_ids_to_associate:
            associate_assets_with_retries(asset_id, asset_ids_to_associate)

        return asset_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_and_wait_until_active_with_retries(event):
    name = event['ResourceProperties']['name']
    model = event['ResourceProperties']['model']
    asset_ids_to_associate = event['ResourceProperties']['asset_ids_to_associate']

    create_asset_response = create_asset_helper(name, model)
    asset_id = create_asset_response['assetId']
    iotsitewise.get_waiter(waiter_name='asset_active').wait(assetId=asset_id)

    # Vladi
    asset_describe = iotsitewise.describe_asset(assetId=asset_id)
    for asset_property in asset_describe['assetProperties']:
        # print(asset_property)
        # print(asset['id'])
        # print(asset_property['id'])
        # print(asset_property['alias'] if 'alias' in asset_property else '')
        iotsitewise.update_asset_property(assetId=asset_id,
                                          propertyId=asset_property['id'],
                                          propertyNotificationState='ENABLED')
    iotsitewise.get_waiter(waiter_name='asset_active').wait(assetId=asset_id)
    # end Vladi

    return asset_id, asset_ids_to_associate

def create_asset_helper(name, model):
    try:
        create_asset_response = iotsitewise.create_asset(assetName=name, assetModelId=model)
    except Exception as e:
        create_asset_response = create_asset_uuid_name_with_retries(name, model)
    return create_asset_response

@retry(wait_fixed=WAIT_TIME_IN_MILLISECONDS, stop_max_attempt_number=STANDARD_RETRY_MAX_ATTEMPT_COUNT)
def create_asset_uuid_name_with_retries(name, model):
    return iotsitewise.create_asset(assetName=name + ' ' + get_timestamp(), assetModelId=model)

@retry(wait_fixed=WAIT_TIME_IN_MILLISECONDS, stop_max_attempt_number=STANDARD_RETRY_MAX_ATTEMPT_COUNT)
def associate_assets_with_retries(asset_id, asset_ids_to_associate):
    desc_response = iotsitewise.describe_asset(assetId=asset_id)
    asset_hierarchies = desc_response['assetHierarchies']
    associate_assets(asset_id, asset_ids_to_associate, asset_hierarchies)

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

@cfnResource.delete
def delete_asset(event, context):
    try:
        asset_id = event['PhysicalResourceId']
        time.sleep(60)

        asset_ids_to_disassociate = event['ResourceProperties']['asset_ids_to_associate']

        if asset_ids_to_disassociate:
            desc_response = iotsitewise.describe_asset(assetId=asset_id)
            asset_hierarchies = desc_response['assetHierarchies']
            disassociate_assets(asset_id, asset_ids_to_disassociate, asset_hierarchies)

        iotsitewise.delete_asset(assetId=asset_id)
        iotsitewise.get_waiter(waiter_name='asset_not_exists').wait(assetId=asset_id)

    except Exception as e:
        logger.error(e)

    finally:
        return event.get('PhysicalResourceId', None)

def associate_assets(asset_id, asset_ids_to_associate, asset_hierarchies):
    try:
        if not asset_ids_to_associate or not asset_hierarchies:
            return True
        hierarchy_id = asset_hierarchies[0]['id']
        for asset_id_to_associate in asset_ids_to_associate:
            iotsitewise.associate_assets(assetId=asset_id, hierarchyId=hierarchy_id, childAssetId=asset_id_to_associate)
        return True
    except Exception as e:
        logger.error(e)
        return False

def disassociate_assets(asset_id, asset_ids_to_disassociate, asset_hierarchies):
    try:
        if not asset_ids_to_disassociate or not asset_hierarchies:
            return True
        hierarchy_id = asset_hierarchies[0]['id']
        for asset_id_to_disassociate in asset_ids_to_disassociate:
            iotsitewise.disassociate_assets(assetId=asset_id, hierarchyId=hierarchy_id, childAssetId=asset_id_to_disassociate)
        return True
    except Exception as e:
        logger.error(e)
        return False

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

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