import sys
import json
import time
import uuid

from kafka import KafkaConsumer, KafkaProducer
import pandas as pd

from graph_based_intrusion_detection.utils import constants
from graph_based_intrusion_detection.utils import config
from graph_based_intrusion_detection.packet_processing.processing_functions import process_packet, process_packets
from graph_based_intrusion_detection.utils.logging import create_logger
from graph_based_intrusion_detection.utils.network_transport_utils import deserialize, serialize


def start_packet_processing(worker_id: str,
                            max_batch_size: int = constants.PACKET_CONSUMER_BATCH_SIZE):
    """
    Worker thread, processes packets from the packet queue and puts the resulting state into the state queue
    :param worker_id: id of the worker
    :param max_batch_size: maximum batch size
    :return: None
    """

    logger = create_logger(f"Packet-Worker {worker_id}")
    logger.info(f"Worker started")

    packet_consumer = KafkaConsumer(constants.PACKET_TOPIC_NAME,
                                    bootstrap_servers='localhost:9092',
                                    auto_offset_reset='earliest')

    state_producer = KafkaProducer(
        bootstrap_servers='localhost:9092', )

    logger.info(f"Initialized, waiting for packets")
    while True:
        # consume batch of packets from kafka
        batch = list()
        while len(batch) < max_batch_size:
            packet_messages = packet_consumer.poll(timeout_ms=constants.PACKET_POLL_TIMEOUT_MS,
                                                   max_records=max_batch_size - len(batch))

            if len(packet_messages) == 0:
                break

            for _, messages in packet_messages.items():
                for message in messages:
                    packet = deserialize(message.value)
                    batch.append(packet)

        if len(batch) == 0:
            time.sleep(0.1)
            continue

        logger.info(f"Processing {len(batch)} packets")

        # process batch
        if len(batch) > 0:
            # process batch with all processing steps
            logger.info(f"Processing {len(batch)} packets in {len(config.PROCESSING_STEPS)} processing steps")

            processing_step_results = dict()
            for processing_step in config.PROCESSING_STEPS:
                current_results = process_packets(batch, processing_step["processing_functions"])
                processing_step_results[processing_step["name"]] = current_results

            serialized_state = serialize(processing_step_results)
            state_producer.send(constants.PROCESSED_STATE_TOPIC_NAME, value=serialized_state)
            logger.info(f"Processed {len(batch)} packets")
        else:
            logger.info(f"No packets to process")
            time.sleep(0.1)


if __name__ == "__main__":
    # get worker id from arguments
    if len(sys.argv) != 2:
        # assign random worker id
        worker_id = uuid.uuid4()
    else:
        worker_id = sys.argv[1]

    start_packet_processing(worker_id)
