import base64
import boto3
import gzip
import uuid
import random

from telemetry import TelemetryData
from telemetry import parse_telemetry_data

telemetry_map = {
    "tidb": [
        "telemetry-tidb-mim.fifo",
        "telemetry-tidb-sync-0.fifo", 
        "telemetry-tidb-sync-1.fifo", 
        "telemetry-tidb-sync-2.fifo",
        "telemetry-tidb-sync-3.fifo",
        "telemetry-tidb-sync-4.fifo",
        "telemetry-tidb-sync-5.fifo",
    ],
}

SQS_MESSAGE_LIMIT = 262144


def lambda_handler(event, context):
    s3 = boto3.client("s3")
    sqs = boto3.client("sqs")

    msg_cnt = 0
    drop_cnt = 0
    for record in event["Records"]:
        bucket = record["s3"]["bucket"]["name"]
        key = record["s3"]["object"]["key"]

        t = key.split("/")[0]
        if not t in telemetry_map:
            continue

        queue = sqs.get_queue_url(QueueName=random.choice(telemetry_map[t]))
        obj = s3.get_object(Bucket=bucket, Key=key)

        with gzip.open(obj["Body"], "rb") as f:
            for line in f.readlines():
                line = line.decode("utf-8")
                telemetry_data: TelemetryData = parse_telemetry_data(line)
                gz_b64 = base64.b64encode(
                    gzip.compress(telemetry_data.__json__().encode("utf-8"))
                ).decode("ascii")
                msg_cnt += 1
                if gz_b64.__sizeof__() >= SQS_MESSAGE_LIMIT:
                    drop_cnt += 1
                else:
                    sqs.send_message(
                        QueueUrl=queue["QueueUrl"],
                        MessageBody=gz_b64,
                        MessageGroupId="telemetry-" + t,
                        MessageDeduplicationId=str(uuid.uuid4()),
                    )
    print("[INSIGHT] msg_cnt={} drop_cnt={}".format(msg_cnt, drop_cnt))


def test_handler(event, context):
    s3 = boto3.client("s3")

    data_list = []
    raw_list = []
    for record in event["Records"]:
        bucket = record["s3"]["bucket"]["name"]
        key = record["s3"]["object"]["key"]

        t = key.split("/")[0]
        if not t in telemetry_map:
            continue

        obj = s3.get_object(Bucket=bucket, Key=key)
        with gzip.open(obj["Body"], "rb") as f:
            for line in f.readlines():
                line = line.decode("utf-8")
                data = parse_telemetry_data(line)
                raw_list.append(line)
                data_list.append(data)

    return raw_list, data_list