# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: MIT-0

import json
import boto3
import time
import os
import logging

logger = logging.getLogger()
logger.setLevel(logging.INFO)

comprehend = boto3.client("comprehend")
sns = boto3.resource("sns")

def lambda_handler(event, context):
    bucketname =  event['Records'][0]['s3']['bucket']['name']
    classifierobject =  event['Records'][0]['s3']['object']['key']
    searchkey='s3://'+ bucketname + '/' + classifierobject
    topic_arn = os.environ.get("TOPIC_ARN")
    sns.Topic(arn=topic_arn).publish(Message='Completed Amazon Comprehend custom classifier training')
    try:
        # Extract Amazon Comprehend document classfier arn
        documentclassifierarn=''
        document_classifier_list = comprehend.list_document_classifiers()
        
        for i in document_classifier_list:
            if(i in 'DocumentClassifierPropertiesList'):
                documentclassifierlist=list(filter(lambda k:k['OutputDataConfig']['S3Uri']==searchkey, document_classifier_list[i]))
                if not documentclassifierlist:
                    documentclassifierarn='Document Classifier ARN not found'
                    logger.info('Document Classifier ARN not found')
                    sns.Topic(arn=topic_arn).publish(Message=documentclassifierarn)
                else:
                    documentclassifierarn=documentclassifierlist[0]['DocumentClassifierArn']
                    logger.info('Document Classifier ARN: ' + documentclassifierarn)
                    message = 'Amazon Comprehend Classifier Model ARN : ' + documentclassifierarn
                    sns.Topic(arn=topic_arn).publish(Message=message)
                    
    except Exception as e: 
            logger.error("Exception (%s)", e)
            raise e

    return {
        'statusCode': 200,
        'EndpointArn': documentclassifierarn
    }