import radix
rtree = radix.Radix()
rnode = rtree.add("192.168.1.0/24")

import pytricia
pyt = pytricia.PyTricia()
pyt.insert('192.168.1.0/24', 'Some value')

import SubnetTree
subnet_tree = SubnetTree.SubnetTree()
for subnet in subnets:
    subnet_tree[subnet] = subnet

results = {}
for ip in ips:
    try:
        matching_subnet = subnet_tree[ip]
        results[ip] = matching_subnet
    except KeyError:
        results[ip] = None

import ipaddress
import os
import pickle
import random
from itertools import combinations

# Path settings
subnets_file = "subnets.pickle"
ips_file = "ips.pickle"

NB_SUBNET = 100000
NB_IPS = 1000
TESTING_ITERATIONS = 3

GENERATE_ALWAYS = True


def generate_subnets():
    subnets = []
    for _ in range(NB_SUBNET):
        prefix_length = random.randint(
            10, 30
        )  # Choose a random prefix length between /10 and /30
        subnet = ipaddress.IPv4Interface(
            f"10.{random.randint(0, 255)}.{random.randint(0, 255)}.{random.randint(0, 255)}/{prefix_length}",
        ).network
        subnets.append(str(subnet))
    return subnets


def generate_ips():
    ips = [
        f"10.{random.randint(0, 255)}.{random.randint(0, 255)}.{random.randint(0, 255)}"
        for _ in range(NB_IPS)
    ]
    return ips


def save_to_pickle(data, file_name):
    with open(file_name, "wb") as f:
        pickle.dump(data, f)


def load_from_pickle(file_name):
    with open(file_name, "rb") as f:
        return pickle.load(f)


# Check for existing datasets and generate if not present
if not os.path.exists(subnets_file) or GENERATE_ALWAYS:
    print("Subnets dataset not found. Generating...")
    subnets = generate_subnets()
    save_to_pickle(subnets, subnets_file)
    print("Subnets dataset generated.")
else:
    subnets = load_from_pickle(subnets_file)
    print("Subnets dataset loaded.")

if not os.path.exists(ips_file) or GENERATE_ALWAYS:
    print("IPs dataset not found. Generating...")
    ips = generate_ips()
    save_to_pickle(ips, ips_file)
    print("IPs dataset generated.")
else:
    ips = load_from_pickle(ips_file)
    print("IPs dataset loaded.")


def match_ips_in_subnets_ipaddress(subnets, ips):
    import ipaddress

    networks = [ipaddress.ip_network(subnet) for subnet in subnets]

    results = {}
    for ip in ips:
        ip_obj = ipaddress.ip_address(ip)
        # Filtering and then directly sorting network objects based on prefixlen
        matching_subnets = sorted(
            [net for net in networks if ip_obj in net],
            key=lambda x: x.prefixlen,
            reverse=True,
        )
        if matching_subnets:
            # The most specific subnet is the one with the largest prefixlength (reverse sorted)
            results[ip] = str(matching_subnets[0])
        else:
            results[ip] = None
    return results


def match_ips_in_subnets_netaddr(subnets, ips):
    from netaddr import IPAddress, IPNetwork

    networks = [IPNetwork(subnet) for subnet in subnets]

    results = {}
    for ip in ips:
        ip_obj = IPAddress(ip)
        matching_subnets = sorted(
            [net for net in networks if ip_obj in net],
            key=lambda x: x.prefixlen,
            reverse=True,
        )
        if matching_subnets:
            # The most specific subnet is the one with the largest prefixlen (reverse sorted)
            results[ip] = str(matching_subnets[0])
        else:
            results[ip] = None
    return results


# def match_ips_in_subnets_netaddripset(subnets, ips):
#     from netaddr import IPNetwork, IPSet

#     # Convert the list of subnet strings to an IPSet of networks
#     subnet_set = IPSet([IPNetwork(subnet) for subnet in subnets])

#     # Convert the list of ip strings to an IPSet of individual addresses
#     ip_set = IPSet(ips)

#     # Find the intersection of IPs and subnets to determine which IPs match which subnets
#     matching_ips = ip_set & subnet_set

#     # `matching_ips` now contains all IPs that fall within the specified subnets
#     # We can now reverse-map these IPs to their respective subnets if needed

#     # Prepare a dictionary to hold the result
#     ip_to_subnet_mapping = {}

#     # For each matching IP, find which subnet it belongs to
#     for ip in matching_ips:
#         for subnet in subnet_set.iter_cidrs():
#             if ip in subnet:
#                 # Convert IP and subnet to strings for easier reading
#                 ip_to_subnet_mapping[str(ip)] = str(subnet)
#                 break  # Stop looking once the correct subnet is found

#     return ip_to_subnet_mapping
# # WILL NOT WORK AS IT WILL FIND THE LARGET MATCHING SUBNET AND NOT THE BEST ONE


def match_ips_in_subnets_pyradix(subnets, ips):
    import radix

    radix_tree = radix.Radix()

    for subnet in subnets:
        rnode = radix_tree.add(subnet)
        rnode.data["subnet"] = subnet

    results = {}
    for ip in ips:
        rnode = radix_tree.search_best(ip)
        if rnode:
            results[ip] = rnode.data["subnet"]
        else:
            results[ip] = None
    return results


def match_ips_in_subnets_pytricia(subnets, ips):
    import pytricia

    pyt = pytricia.PyTricia()

    for subnet in subnets:
        pyt[subnet] = subnet

    results = {}
    for ip in ips:
        try:
            most_specific_subnet = pyt[ip]
            results[ip] = most_specific_subnet
        except KeyError:
            results[ip] = None
    return results


def match_ips_in_subnets_PySubnetTree(subnets, ips):
    import SubnetTree

    subnet_tree = SubnetTree.SubnetTree()

    # Fill the SubnetTree with subnets. You can associate data with each subnet if needed.
    for subnet in subnets:
        subnet_tree[subnet] = subnet

    results = {}
    for ip in ips:
        try:
            # The lookup will succeed if the IP is in any of the subnets added to the tree.
            # It returns the data associated with the subnet, which in this case is the subnet itself.
            matching_subnet = subnet_tree[ip]
            results[ip] = matching_subnet
        except KeyError:
            # If the IP does not match any subnet, catch the exception and set the result to None.
            results[ip] = None

    return results


def ip_to_binary(ip):
    return "".join([bin(int(x) + 256)[3:] for x in ip.split(".")])


def ip_to_int(ip):
    """Convert an IP string to an integer."""
    return sum(int(byte) << 8 * i for i, byte in enumerate(reversed(ip.split("."))))


def match_ips_in_subnets_binary(subnets, ips):
    """Match IPs to the most specific subnet."""
    results = {}

    subnet_info = []
    for subnet in subnets:
        subnet_addr, prefix_len = subnet.split("/")
        subnet_int = ip_to_int(subnet_addr)
        mask = (-1 << (32 - int(prefix_len))) & 0xFFFFFFFF
        subnet_info.append((subnet_int & mask, mask, subnet))

    for ip in ips:
        ip_int = ip_to_int(ip)
        best_match = None
        for subnet_int, mask, subnet in subnet_info:
            if ip_int & mask == subnet_int:
                if best_match is None or best_match[1] < mask:
                    best_match = (subnet, mask)
        results[ip] = best_match[0] if best_match else None

    return results


import timeit

import matplotlib.pyplot as plt


def dynamic_timing_and_plotting(subnets, ips):
    execution_times = {}
    for name, func in globals().items():
        if callable(func) and name.startswith("match_ips_in_subnets_"):
            tool_name = name.split("match_ips_in_subnets_")[-1]
            time_taken = timeit.timeit(
                lambda: func(subnets, ips), number=TESTING_ITERATIONS
            )
            execution_times[tool_name] = time_taken

    from pprint import pprint

    pprint(execution_times)

    # Plotting the execution times
    plt.figure(figsize=(10, 6))
    plt.bar(execution_times.keys(), execution_times.values(), color="skyblue")
    plt.xlabel("Tool")
    plt.ylabel("Time (seconds)")
    plt.title("Execution Time for Matching IPs in Subnets")
    plt.xticks(rotation=45)
    plt.tight_layout()
    plt.show()


def testing_tools():
    results = {}
    for name, func in globals().items():
        if callable(func) and name.startswith("match_ips_in_subnets_"):
            tool_name = name.split("match_ips_in_subnets_")[-1]
            results[tool_name] = func(subnets, ips)

    discrepancies = []
    matches = []

    # Compare all pairs of results for discrepancies
    for (tool1, result1), (tool2, result2) in combinations(results.items(), 2):
        if result1 != result2:
            discrepancies.append((tool1, tool2))
        else:
            matches.append((tool1, tool2))

    # Printing the outcome summary
    if not discrepancies:
        print("All tools have produced identical results.")
    else:
        print("Discrepancies found between the following tools:")
        for tool1, tool2 in discrepancies:
            print(f"- {tool1} vs {tool2}")

    # Also summarize tools that had matching results for completeness
    if matches:
        print("\nTools with matching results:")
        for tool1, tool2 in matches:
            print(f"- {tool1} and {tool2} have identical results")

testing_tools()

dynamic_timing_and_plotting(subnets, ips)