import sys
import uuid
import time
from datetime import datetime

from kafka import KafkaConsumer, KafkaProducer
import redis

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


def start_work_dispatcher(worker_id: str):
    """
     Worker function, creates node registry based on network state and dispatches work based on nodes that need updating
    :param worker_id: id of the worker
    :return: None
    """

    logger = create_logger(f"Work-Dispatcher {worker_id}")

    # connect to kafka
    node_update_producer = KafkaProducer(
        bootstrap_servers='localhost:9092')

    node_analysis_results_consumer = KafkaConsumer(
        constants.NODE_ANALYSIS_RESULTS_TOPIC,
        bootstrap_servers='localhost:9092',
        auto_offset_reset='earliest',
    )
    logger.info("Connected to kafka")

    # connect to redis
    redis_host = redis.Redis(host='localhost', port=6379, db=0)
    logger.info("Connected to redis")

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

    awaiting_results = list()

    # define last values for better logging
    last_network_size_key = "last_network_size"
    last_num_awaiting_results_key = "last_num_awaiting_results"
    last_values = {
        last_network_size_key: 0,
        last_num_awaiting_results_key: 0
    }

    while True:
        # get current network state
        network_state_raw = redis_host.get(constants.NETWORK_STATE_KEY)

        if network_state_raw is None:
            logger.info("No network state found, waiting for state")
            time.sleep(1)
            continue

        network_state = deserialize(network_state_raw)

        # get nodes from network state
        nodes = list(network_state[constants.NETWORK_STATE_CONNECTION_GRAPH_KEY].nodes())
        if len(nodes) != last_values[last_network_size_key]:
            logger.info(f"Found {len(nodes)} nodes in network state")
            last_values[last_network_size_key] = len(nodes)

        # obtain registry lock
        lock = redis_host.lock(constants.NODE_REGISTRY_LOCK, timeout=1000)
        if lock.acquire(blocking=False):
            try:
                logger.debug("Acquired node registry lock")

                # get current node registry
                node_registry_raw = redis_host.get(constants.NODE_REGISTRY_KEY)
                if node_registry_raw is None:
                    node_registry = {}
                else:
                    node_registry = deserialize(node_registry_raw)

                # remove nodes from registry that are not in the network state
                for node in list(node_registry.keys()):
                    if node not in nodes:
                        del node_registry[node]

                # get node analysis results since last update
                node_analysis_results = list()
                while True:
                    node_analysis_results_raw = node_analysis_results_consumer.poll(
                        timeout_ms=constants.NODE_ANALYSIS_RESULTS_POLL_TIMEOUT_MS)

                    if len(node_analysis_results_raw) == 0:
                        break

                    for topic, messages in node_analysis_results_raw.items():
                        for message in messages:
                            result_dict = deserialize(message.value)
                            node = result_dict["node"]
                            node_analysis_results.append(result_dict)

                            # remove results from awaiting results
                            if node in awaiting_results:
                                awaiting_results.remove(node)

                if len(awaiting_results) != last_values[last_num_awaiting_results_key]:
                    logger.info(f"Found {len(awaiting_results)} nodes with awaiting results")
                    last_values[last_num_awaiting_results_key] = len(awaiting_results)

                # process node analysis results
                for analysis_result in node_analysis_results:
                    node = analysis_result["node"]
                    update_time = analysis_result["time"]
                    results = analysis_result["results"]

                    # update node states in registry based on results
                    if node in node_registry:
                        node_registry[node]["last_updated"] = update_time
                        node_registry[node]["results"] = results
                    else:
                        logger.warning(f"Received result for unknown node {node}")

                # get nodes that are not in the registry
                new_nodes = [node for node in nodes if node not in node_registry]

                # add new nodes to the registry
                if len(new_nodes) > 0:
                    logger.info(f"Found {len(new_nodes)} new nodes in network state")
                for node in new_nodes:
                    node_registry[node] = {"last_updated": None}

                # create and dispatch work packages
                new_node_updates = list()
                for node, node_data in node_registry.items():
                    last_updated = node_data["last_updated"]
                    if last_updated is None or (time.time() - last_updated) > (
                            constants.NODE_ANALYSIS_INTERVAL_MS / 1000):
                        # only dispatch work if the node result is not already being awaited
                        if node not in awaiting_results:
                            node_update_message = {
                                "node": node,
                                "time": time.time()
                            }
                            node_update_producer.send(constants.NODE_ANALYSIS_TOPIC,
                                                      value=serialize(node_update_message))
                            new_node_updates.append(node)
                            awaiting_results.append(node)

                if len(new_node_updates) > 0:
                    logger.info(f"Dispatched {len(new_node_updates)} new node updates")

                # update node registry
                redis_host.set(constants.NODE_REGISTRY_KEY, serialize(node_registry))
            finally:
                # release lock
                lock.release()
                logger.debug("Released node registry lock")


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

    start_work_dispatcher(worker_id)
