import networkx as nx
import pandas as pd
import numpy as np


def extract_to_graph(input_graph: nx.Graph,
                     packet: dict | pd.Series,
                     source_field: str,
                     destination_field: str,
                     per_packet_fields: list[str],
                     per_second_fields: list[str],
                     timedelta_field: str):
    """
    Extract the given fields from the packet and add them to the graph
    :param input_graph: graph to update
    :param packet: packet to analyze
    :param source_field: source field
    :param destination_field: destination field
    :param per_packet_fields: fields to extract per packet
    :param per_second_fields: fields to extract per second
    :param timedelta_field: field containing the timedelta to the start of the capture
    :return: updated graph
    """
    # get timedelta to start of capture
    timedelta = packet[timedelta_field]

    # get current second
    try:
        second = int(timedelta.total_seconds())
    except ValueError:
        # unusable packet
        return input_graph

    # get src and destination address
    src = packet[source_field]
    dst = packet[destination_field]

    if pd.isna(src) or pd.isna(dst):
        # unusable packet
        return input_graph

    # add nodes, if not already present
    if not input_graph.has_node(src):
        input_graph.add_node(src)
    if not input_graph.has_node(dst):
        input_graph.add_node(dst)

    # check content of packet to determine layer
    if packet["tcp/tcp.srcport"] is not np.nan:
        min_layer = 4
    elif packet["ip/ip.src"] is not np.nan:
        min_layer = 3
    elif packet["eth/eth.src"] is not np.nan:
        min_layer = 2
    else:
        min_layer = 1

    # add edge, if not already present
    if not input_graph.has_edge(src, dst):
        input_graph.add_edge(src, dst)

        # add per packet fields
        for field in per_packet_fields:
            if field["min_layer"] <= min_layer:
                value = field["base_value_function"](packet)
                if value is not None:
                    input_graph[src][dst][field["name"]] = value

        # add per second fields
        for field in per_second_fields:
            if field["min_layer"] <= min_layer:
                value = field["base_value_function"](packet)
                if value is not None:
                    input_graph[src][dst][field["name"]] = {second: value}

    else:

        # update per packet fields
        for field in per_packet_fields:
            if field["min_layer"] <= min_layer:
                value = field["update_function"](input_graph[src][dst][field["name"]], packet)
                if value is not None:
                    input_graph[src][dst][field["name"]] = value

        # update per second fields
        for field in per_second_fields:
            if field["name"] in input_graph[src][dst]:
                if second not in input_graph[src][dst][field["name"]]:
                    if field["min_layer"] <= min_layer:
                        value = field["base_value_function"](packet)
                        if value is not None:
                            input_graph[src][dst][field["name"]][second] = value
                else:
                    if field["min_layer"] <= min_layer:
                        value = field["update_function"](input_graph[src][dst][field["name"]][second], packet)
                        if value is not None:
                            input_graph[src][dst][field["name"]][second] = value
            else:
                if field["min_layer"] <= min_layer:
                    value = field["base_value_function"](packet)
                    if value is not None:
                        input_graph[src][dst][field["name"]] = {second: value}

    return input_graph


def update_layer_3_graph(packet: pd.Series,
                         state: dict,
                         per_packet_fields: list,
                         per_second_fields: list,
                         timedelta_field: str) -> nx.Graph:
    """
    Update the layer 3 graph
    :param packet: packet to analyze
    :param state: current state of the model
    :param per_packet_fields: fields to extract per packet
    :param per_second_fields: fields to extract per second
    :param timedelta_field: field containing the timedelta to the start of the capture
    :return: updated state
    """
    layer_3_graph = state if state is not None else nx.DiGraph()

    layer_3_graph = extract_to_graph(layer_3_graph, packet,
                                     "ip/ip.src",
                                     "ip/ip.dst",
                                     per_packet_fields,
                                     per_second_fields,
                                     timedelta_field)

    return layer_3_graph


def update_layer_2_graph(packet: pd.Series,
                         state: dict,
                         per_packet_fields: list,
                         per_second_fields: list,
                         timedelta_field: str) -> nx.Graph:
    """
    Update the layer 2 graph
    :param packet: packet to analyze
    :param state: current state of the model
    :param per_packet_fields: fields to extract per packet
    :param per_second_fields: fields to extract per second
    :param timedelta_field: field containing the timedelta to the start of the capture
    :return: updated state
    """
    layer_2_graph = state if state is not None else nx.DiGraph()

    layer_2_graph = extract_to_graph(layer_2_graph, packet,
                                     "eth/eth.src",
                                     "eth/eth.dst",
                                     per_packet_fields,
                                     per_second_fields,
                                     timedelta_field)

    return layer_2_graph
