import numpy as np

# define processing functions
from functools import partial

from graph_based_intrusion_detection.packet_processing.packet_inspection import update_port_list
from graph_based_intrusion_detection.utils import constants

from graph_based_intrusion_detection.graph_processing.data_extraction import create_dataset_for_node

from graph_based_intrusion_detection.packet_processing.graph_processing import update_layer_3_graph
from graph_based_intrusion_detection.utils.state_merging import merge_graph
from graph_based_intrusion_detection.analysis.result_interpretation import interpret_prediction_result

PER_CONNECTION_FIELDS = [
    {
        "name": "total_packet_count",
        "min_layer": 2,
        "base_value_function": lambda _: 1,
        "update_function": lambda x, _: x + 1
    },
    {
        "name": "total_bytes",
        "min_layer": 2,
        "base_value_function": lambda packet: packet["frame/frame.len"],
        "update_function": lambda x, packet: x + packet["frame/frame.len"]
    },
    {
        "name": "src_ports",
        "min_layer": 4,
        "base_value_function": lambda packet: update_port_list(packet, [], "tcp/tcp.srcport"),
        "update_function": lambda x, packet: update_port_list(packet, x, "tcp/tcp.srcport"),
    },
    {
        "name": "dst_ports",
        "min_layer": 4,
        "base_value_function": lambda packet: update_port_list(packet, [], "tcp/tcp.dstport"),
        "update_function": lambda x, packet: update_port_list(packet, x, "tcp/tcp.dstport"),
    },
]

PER_SECOND_FIELDS = [
    {
        "name": "packets_per_second",
        "min_layer": 2,
        "base_value_function": lambda _: 1,
        "update_function": lambda x, _: x + 1
    },
    {
        "name": "bytes_per_second",
        "min_layer": 2,
        "base_value_function": lambda packet: packet["frame/frame.len"],
        "update_function": lambda x, packet: x + packet["frame/frame.len"]
    },
    {
        "name": "max_tcp_segment_size",
        "min_layer": 4,
        "base_value_function": lambda packet: packet["tcp/tcp.options/tcp.options.mss/tcp.options.mss_val"],
        "update_function": lambda x, packet:
        packet["tcp/tcp.options/tcp.options.mss/tcp.options.mss_val"] if x < packet[
            "tcp/tcp.options/tcp.options.mss/tcp.options.mss_val"] else x
    },
    {
        "name": "max_tcp_header_length",
        "min_layer": 4,
        "base_value_function": lambda packet: packet["tcp/tcp.hdr_len"],
        "update_function": lambda x, packet:
        packet["tcp/tcp.hdr_len"] if x < packet["tcp/tcp.hdr_len"] else x
    },
    {
        "name": "max_frame_length",
        "min_layer": 2,
        "base_value_function": lambda packet: packet["frame/frame.len"],
        "update_function": lambda x, packet: packet["frame/frame.len"] if x < packet["frame/frame.len"] else x
    },
    {
        "name": "min_frame_length",
        "min_layer": 2,
        "base_value_function": lambda packet: packet["frame/frame.len"],
        "update_function": lambda x, packet: packet["frame/frame.len"] if x > packet["frame/frame.len"] else x
    }
]

PROCESSING_STEPS = [
    {
        "name": "layer_3_graph",
        "processing_functions": [
            partial(update_layer_3_graph, per_packet_fields=PER_CONNECTION_FIELDS, per_second_fields=PER_SECOND_FIELDS,
                    timedelta_field=constants.TIMEDELTA_FIELD)
        ],
        "merge_function": merge_graph
    }
]

# define fields for connection analysis
NODE_ANALYSIS_PER_SECOND_ATTRIBUTES = [x["name"] for x in PER_SECOND_FIELDS]
NODE_ANALYSIS_PER_CONNECTION_ATTRIBUTES = [x["name"] for x in PER_CONNECTION_FIELDS]

# define analysis model and scaler paths
ANALYSIS_STEPS = [
    {
        "name": "incoming_attack_probability",
        "network_state_key": "layer_3_graph",
        "dataset_function": partial(create_dataset_for_node,
                                    per_second_attributes=NODE_ANALYSIS_PER_SECOND_ATTRIBUTES,
                                    per_connection_attributes=NODE_ANALYSIS_PER_CONNECTION_ATTRIBUTES),
        "dataset_identifiers": ["node"],
        "model_path": "/home/alex/projects/it-security-praktikum/devices/dev08-echo-dot-l4s3re/graph_based_intrusion_detection/resources/models/"
                      "incoming_attack_clf.pickle",
        "scaler_path": "/home/alex/projects/it-security-praktikum/devices/dev08-echo-dot-l4s3re/graph_based_intrusion_detection/resources/models/"
                       "layer_3_scaler.pickle",
        "predict_function": "predict_proba",
        "interpretation_function": lambda x: np.max(x[:, 1])
    },
    {
        "name": "outgoing_attack_probability",
        "network_state_key": "layer_3_graph",
        "dataset_function": partial(create_dataset_for_node,
                                    per_second_attributes=NODE_ANALYSIS_PER_SECOND_ATTRIBUTES,
                                    per_connection_attributes=NODE_ANALYSIS_PER_CONNECTION_ATTRIBUTES),
        "dataset_identifiers": ["node"],
        "model_path": "/home/alex/projects/it-security-praktikum/devices/dev08-echo-dot-l4s3re/graph_based_intrusion_detection/resources/models/"
                      "outgoing_attack_nn_clf.pickle",
        "scaler_path": "/home/alex/projects/it-security-praktikum/devices/dev08-echo-dot-l4s3re/graph_based_intrusion_detection/resources/models/"
                       "layer_3_scaler.pickle",
        "predict_function": "predict",
        "interpretation_function": lambda x: np.max(x[:, 1])
    }

]
