import sys
import uuid

import pandas as pd
from kafka import KafkaConsumer
import redis

from graph_based_intrusion_detection.utils import constants
from graph_based_intrusion_detection.utils import config
from graph_based_intrusion_detection.utils.logging import create_logger
from graph_based_intrusion_detection.utils.state_merging import merge_states
from graph_based_intrusion_detection.utils.network_transport_utils import deserialize, serialize


def start_state_merging(worker_id: str):
    """
    Worker function, processes states from the state queue and merges them, updating the state in the redis database
    :param worker_id:
    :return: None
    """

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

    # establish connection to kafka
    state_consumer = KafkaConsumer(constants.PROCESSED_STATE_TOPIC_NAME,
                                   bootstrap_servers='localhost:9092',
                                   auto_offset_reset='earliest')

    logger.info("Connected to kafka")

    # establish connection to redis
    redis_host = redis.Redis(host='localhost', port=6379, db=0)

    logger.info("Connected to redis")

    logger.info("Initialized, waiting for states")

    while True:
        # consume a state from kafka
        batch = list()
        while len(batch) < constants.STATE_CONSUMER_BATCH_SIZE:
            state = state_consumer.poll(timeout_ms=constants.STATE_POLL_TIMEOUT_MS,
                                        max_records=constants.STATE_CONSUMER_BATCH_SIZE - len(batch))
            if len(state) == 0:
                break

            for _, records in state.items():
                for record in records:
                    state = deserialize(record.value)
                    batch.append(state)

        if len(batch) == 0:
            continue

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

        # acquire lock
        lock = redis_host.lock(constants.STATE_LOCK_NAME, timeout=5)

        # merge states in batch before updating the state in redis to keep lock time short
        batch_states_merged = dict()
        for processing_step in config.PROCESSING_STEPS:
            step_name = processing_step["name"]
            merged_step = batch[0][step_name]
            for i in range(len(batch) - 1):
                merged_step = processing_step["merge_function"](batch[i][step_name], batch[i + 1][step_name])

            batch_states_merged[step_name] = merged_step

        if lock.acquire(blocking=False):
            try:
                logger.debug("Acquired lock")
                # get current state from redis
                current_state_raw = redis_host.get(constants.NETWORK_STATE_KEY)

                if current_state_raw is not None:
                    current_state = deserialize(current_state_raw)
                else:
                    current_state = dict()

                # merge states
                for step in config.PROCESSING_STEPS:
                    # only merge if step is present in batch, otherwise just use new state
                    step_name = step["name"]
                    if step_name not in current_state:
                        current_state[step_name] = batch_states_merged[step_name]
                    else:
                        current_state[step_name] = step["merge_function"](current_state[step_name],
                                                                          batch_states_merged[step_name])

                serialized_state = serialize(current_state)
                redis_host.set(constants.NETWORK_STATE_KEY, serialized_state)
            except Exception as e:
                logger.error(f"Error while merging states: {e}")
            finally:
                logger.debug("Releasing lock")
                lock.release()

            logger.info(f"Merged {len(batch)} state(s), waiting for next batch")
        else:
            logger.info("Could not acquire lock, retrying in 1 second")
            continue


if __name__ == "__main__":
    if len(sys.argv) < 2:
        worker_id = str(uuid.uuid4())
    else:
        worker_id = sys.argv[1]

    start_state_merging(worker_id)
