import base64
import json
from datetime import datetime
import boto3

DEBUG_PRINT = False
MAX_BATCH_SIZE = 150

kinesis_client = boto3.client("kinesis")
cloudwatch_client = boto3.client("cloudwatch")

whilte_list_event_name = [
    "Audio-Player-Event",
    "Audio-Creation-Event",
    "Audio-Action-Event",
    "Playlist-Action-Event",
    "Link-Click-Event",
    "Web-Page-Event",
    "Artist-Action-Event",
    "General-Event",
    "Audio-Editing-Event",
    "Account-Action-Event",
    "Search-Event",
    "App-Audio-Player-Event",
    "App-Audio-Creation-Event",
    "App-Audio-Action-Event",
    "App-Playlist-Action-Event",
    "App-Link-Click-Event",
    "App-Artist-Action-Event",
    "App-General-Event",
    "App-Audio-Editing-Event",
    "App-Account-Action-Event",
    "App-Search-Event",
    "Web-User-Event",
    "Hook-Web-Event",
]


def debug_print(message):
    if DEBUG_PRINT:
        print(message)


def is_whitelisted_event(json_value):
    return json_value["event"] in whilte_list_event_name


def notify_cloudwatch_oversized_batch(
    stream_name, client_ip, request_time, domain_name, batched_data
):
    """
    Notifies CloudWatch that a batch of data is too large to be processed.
    """
    debug_print(
        f"Notifying CloudWatch of oversized batch with {len(batched_data)} events..."
    )

    # puts the batch size into CloudWatch by stream name
    cloudwatch_client.put_metric_data(
        Namespace="EventLogger/FilterLambda",
        MetricData=[
            {
                "MetricName": "BatchExceedsMaxCount",
                "Dimensions": [
                    {"Name": "StreamName", "Value": stream_name},
                ],
                "Value": len(batched_data),
                "Unit": "Count",
                "Timestamp": datetime.now(),
            }
        ],
    )

    debug_print("Sent event to CloudWatch Metrics.")
    return


def _base64_encode_json(json_data):
    """
    Returns a base64-encoded string of the given JSON data.
    """
    return base64.b64encode(json.dumps(json_data).encode("utf-8")).decode("utf-8")


def handle_batch(stream_name, client_ip, request_time, domain_name, batched_data):
    """
    Handles a batch of data that needs to be sent along in the event pipeline to Kinesis.
    The data should consist of a list of events, where each event is a JSON object.
    Each event is base64-encoded and then sent to Kinesis along with the client IP,
    request time, and domain name associated with the request it originated from. The
    partition key is the value of the anonymous ID taken from the event. amogus
    """
    records = []

    # process each data record in the batch into the format expected by Kinesis
    for event in batched_data:
        if not is_whitelisted_event(event):
            debug_print(
                f"Failed to recognize event type: {event['event']}. Dropping event: {event}"
            )
            continue

        request_data = {
            "request_body": _base64_encode_json(event),
            "x-client-ip": client_ip,
            "request-time": request_time,
            "domain-name": domain_name,
        }
        record = {
            "Data": json.dumps(request_data),
            "PartitionKey": event["anonymousId"],
        }
        records.append(record)
        debug_print(f"Added record to batch: {record}")

    # send the batch to Kinesis
    kinesis_client.put_records(
        StreamName=stream_name,
        Records=records,
    )
    return


def lambda_handler(event, context):
    debug_print("Starting pre-Kinesis data filtering lambda...")

    # perform some data validation on the event
    assert "StreamName" in event, ""
    assert "ClientIP" in event, ""
    assert "RequestTime" in event, ""
    assert "DomainName" in event, ""
    assert "Data" in event, ""
    assert isinstance(event["StreamName"], str), ""
    assert isinstance(event["ClientIP"], str), ""
    assert isinstance(event["RequestTime"], str), ""
    assert isinstance(event["DomainName"], str), ""
    assert isinstance(event["Data"], list), ""

    stream_name, client_ip, request_time, domain_name, data = (
        event["StreamName"],
        event["ClientIP"],
        event["RequestTime"],
        event["DomainName"],
        event["Data"],
    )

    debug_print(f"Received batch of {len(data)} events.")
    debug_print(
        f"Stream name: {stream_name}\nClient IP: {client_ip}\nRequest time: {request_time}\nDomain name: {domain_name}\nData: {data}"
    )

    # if the batch size exceeds the maximum, send a notification to CloudWatch metrics
    if len(data) > MAX_BATCH_SIZE:
        debug_print(f"Batch size {len(data)} exceeds limit of {MAX_BATCH_SIZE}.")
        notify_cloudwatch_oversized_batch(
            stream_name, client_ip, request_time, domain_name, data
        )

    debug_print("Sending batch to Kinesis...")

    # send the batch to Kinesis
    handle_batch(stream_name, client_ip, request_time, domain_name, data)

    debug_print("Sent batch to Kinesis.")

    debug_print("Lambda finished executing.")