import ijson as ijson

from mona_sdk import MonaSingleMessage, Client

# Set the path of the JSON file and the num of objects be batch
PATH_TO_RAW_RATINGS = "/home/smadar/workspace/mona-sdk/mona_sdk/ratings.json"
NUMBER_OF_MESSAGES_IN_BATCH = 100000

# Set the client
sdk_client = Client(
    api_key="75425b87-26df-4fce-948e-4e64540d3d8e",
    secret="e77a979d-8362-4f4b-812b-5552d223e540",
    user_id="69d85e92-2894-4f7c-a1fb-55312711745f",
    raise_authentication_exceptions=True,
    raise_export_exceptions=True,
)

messages = []
number_of_messages_left_in_current_batch = NUMBER_OF_MESSAGES_IN_BATCH
number_of_batches = 0
total_number_of_messages_sent = 0
number_of_items_iterated = 0  # Where the script is in the file, how many messages
# it already went through (can be different from total_number_of_messages_sent when
# there are failures).


# Set the record keys
with open(PATH_TO_RAW_RATINGS, "r") as read_fd:
    for record in ijson.items(read_fd, "item"):
        messages.append(
            MonaSingleMessage(
                message=record,
                contextClass="ratings",
                contextId=record["user"] + "_" + record["movie"],
                exportTimestamp=record["timestamp"],
            )
        )
        number_of_messages_left_in_current_batch -= 1

        # If we processed the number we wanted in a batch, then it's time to send
        # this batch.
        if number_of_messages_left_in_current_batch == 0:
            number_of_items_iterated += NUMBER_OF_MESSAGES_IN_BATCH
            try:
                # sdk_client.export_batch(messages)
                number_of_batches += 1
                total_number_of_messages_sent += NUMBER_OF_MESSAGES_IN_BATCH
                print(
                    f"Successfully sent messages: "
                    f"{number_of_items_iterated - NUMBER_OF_MESSAGES_IN_BATCH} to "
                    f"{number_of_items_iterated}"
                )
            except Exception:
                print(
                    f"Failed to send messages: "
                    f"{number_of_items_iterated - NUMBER_OF_MESSAGES_IN_BATCH} to "
                    f"{number_of_items_iterated}"
                )

            messages = []
            number_of_messages_left_in_current_batch = NUMBER_OF_MESSAGES_IN_BATCH

    # Here we process the leftovers that didn't reach the amount of the full batch.
    number_of_items_iterated += len(messages)
    try:
        # sdk_client.export_batch(messages)
        number_of_batches += 1
        total_number_of_messages_sent += len(messages)
        print(
            f"Sent msgs {number_of_items_iterated - len(messages)} to "
            f"{number_of_items_iterated}"
        )
    except Exception:
        print(
            f"Failed to send messages {number_of_items_iterated - len(messages)} to "
            f"{number_of_items_iterated}"
        )
    print(
        f"Sent {total_number_of_messages_sent} messages in {number_of_batches} batches."
    )
