import os
import sys
import time
import boto3
from botocore.exceptions import ClientError
from datetime import datetime, timedelta
from kinesis_video_streams_parser import KvsConsumerLibrary
from kinesis_video_fragment_processor import KvsFragementProcessor

rekognition_client = boto3.client("rekognition")
kvs_client = boto3.client('kinesisvideo')
kvs_fragment_processor = KvsFragementProcessor()
last_good_fragment_tags = None

stream_name = "" ## Provide KVS name here (stream must already exist)
stream_arn = "" ## Provide KVS ARN here (stream must already exist)
stream_processor_name = "" ## Provide a name for the stream processor
s3_bucket_name = "" ## Provide the S3 bucket
role_arn = "" ## Provide the role ARN for Rekognition
sns_topic_arn = "" ## Provide the SNS ARN


#############################################
## KVS Consumer Library Callbacks

def on_fragment_arrived(stream_name, fragment_bytes, fragment_dom, fragment_receive_duration):
        try:
            last_good_fragment_tags = kvs_fragment_processor.get_fragment_tags(fragment_dom)
            fragment_num = last_good_fragment_tags['AWS_KINESISVIDEO_FRAGMENT_NUMBER']
            rekognition_client.start_stream_processor(
                Name=stream_processor_name,
                StartSelector={
                    'KVSStreamStartSelector': {
                        'FragmentNumber': fragment_num,
                    }
                },
                StopSelector={
                    'MaxDurationInSeconds': 2
                }
            )
            time.sleep(2)
        except Exception as err:
            print(err)

def on_stream_read_complete(stream_name):
    print(f"Stream {stream_name} read complete")

def on_stream_read_exception(stream_name, error):
    print(f"Stream {stream_name} read exception: {error}")

## Main Lambda Function

def lambda_handler(event, context):
    ## Step 1: Ensure stream processor is deleted
    try:
        rekognition_client.delete_stream_processor(
            Name=stream_processor_name,
        )
        rekognition_client.create_stream_processor(
            Input={
                'KinesisVideoStream': {
                    'Arn': stream_arn
                }
            },
            Output={
                'S3Destination': {
                    'Bucket': s3_bucket_name,
                    'KeyPrefix': 'stream-results'
                }
            },
            Name=stream_processor_name,
            Settings={
                'ConnectedHome': {
                    'Labels': [
                        'PERSON',
                    ],
                    'MinConfidence': 80
                }
            },
            RoleArn=role_arn,
            NotificationChannel={
                'SNSTopicArn':sns_topic_arn
            }
        )
    except:
    ## Step 2: Create Rekognition Stream Processor
        rekognition_client.create_stream_processor(
            Input={
                'KinesisVideoStream': {
                    'Arn': stream_arn
                }
            },
            Output={
                'S3Destination': {
                    'Bucket': s3_bucket_name,
                    'KeyPrefix': 'stream-results'
                }
            },
            Name=stream_processor_name,
            Settings={
                'ConnectedHome': {
                    'Labels': [
                        'PERSON',
                    ],
                    'MinConfidence': 80
                }
            },
            RoleArn=role_arn,
            NotificationChannel={
                'SNSTopicArn':sns_topic_arn
            }
        )

    ## Step 3: Prepare connection to stream
    kvs_client = boto3.client('kinesisvideo')

    response = kvs_client.get_data_endpoint(
        StreamARN=stream_arn,
        APIName='GET_MEDIA'
    )

    get_endpoint = response['DataEndpoint']

    kvs_media_client = boto3.client('kinesis-video-media', endpoint_url=get_endpoint)
    get_media_response = kvs_media_client.get_media(
        StreamName=stream_name,
        StartSelector={
            'StartSelectorType': 'NOW'
        }
    )

    ## Step 4: Prepare consumer library
    my_stream01_consumer = KvsConsumerLibrary(stream_name, 
                                              get_media_response, 
                                              on_fragment_arrived, 
                                              on_stream_read_complete, 
                                              on_stream_read_exception
                                            )
    
    ## Step 5: Run consumer
    my_stream01_consumer.run()