import sys
import joblib
import uuid
import time

from kafka import KafkaConsumer, KafkaProducer
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.network_transport_utils import deserialize, serialize
from graph_based_intrusion_detection.graph_processing.data_extraction import create_dataset_for_node
from graph_based_intrusion_detection.analysis.analysis_steps import run_analysis_step


def start_connection_analysis(worker_id: str) -> None:
    """
    Starts the connection analysis worker
    :param worker_id: id of the worker
    :return: None
    """

    logger = create_logger(f"Connection-Analysis {worker_id}")
    logger.info(f"Connection-Analysis {worker_id} started")

    # load analysis models
    logger.info(f"Loading analysis models")
    analysis_models = config.ANALYSIS_STEPS
    for step in analysis_models:
        # load model and scaler for analysis step
        scaler = joblib.load(step["scaler_path"])
        model = joblib.load(step["model_path"])
        prediction_function = eval(f"model.{step['predict_function']}")

        # append model and scaler to step
        step["scaler"] = scaler
        step["model"] = model
        step["predict_function"] = prediction_function

        logger.info(f"Loaded model {step['name']}")

    # connect to connection update queue
    node_update_consumer = KafkaConsumer(
        constants.NODE_ANALYSIS_TOPIC,
        bootstrap_servers='localhost:9092',
        auto_offset_reset='earliest')
    node_update_result_producer = KafkaProducer(
        bootstrap_servers="localhost:9092",
    )

    # connect to redis
    redis_client = redis.StrictRedis(host="localhost", port=6379, db=0)

    while True:
        # get batch of connections from kafka to analyze
        node_batch = list()
        while len(node_batch) < constants.NODE_ANALYSIS_BATCH_SIZE:
            node_messages = node_update_consumer.poll(timeout_ms=constants.NODE_ANALYSIS_POLL_TIMEOUT_MS,
                                                      max_records=constants.NODE_ANALYSIS_BATCH_SIZE - len(
                                                          node_batch))

            if len(node_messages) == 0:
                break

            for _, messages in node_messages.items():
                for message in messages:
                    node = deserialize(message.value)
                    node_batch.append(node)

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

        # remove update message older than the timeout
        node_batch = [node["node"] for node in node_batch if
                      (time.time() - node["time"]) < constants.NODE_ANALYSIS_TIMEOUT_MS / 1000]

        if len(node_batch) == 0:
            continue

        # get network state from redis, dirty read is not really a problem since updates will be processed in the
        # next iteration
        network_state_raw = redis_client.get(constants.NETWORK_STATE_KEY)

        if network_state_raw is None:
            logger.critical(f"Network state not found in redis, retrying on next iteration")
            time.sleep(1)
            continue

        network_state = deserialize(network_state_raw)

        # run all analysis steps
        logger.info(f"Analyzing {len(node_batch)} nodes")
        step_results = {step["name"]: run_analysis_step(node_batch, step, network_state) for step in analysis_models}

        # compose result dict for each node and send result to kafka topic
        for node in node_batch:
            node_results = {step_name: step_results[node] for step_name, step_results in step_results.items()}

            # compose message with node, time and results
            result_message = {
                "node": node,
                "time": time.time(),
                "results": node_results
            }

            # send result to kafka
            node_update_result_producer.send(constants.NODE_ANALYSIS_RESULTS_TOPIC, serialize(result_message))


if __name__ == "__main__":
    if len(sys.argv) < 2:
        worker_id = str(uuid.uuid4())
        print(f"Worker id not provided, using {worker_id}")
    else:
        worker_id = sys.argv[1]

    start_connection_analysis(worker_id)
