from ipaddress import IPv4Network, IPv6Network

from csirtmu.beta_temporal.external_services.alpha.structs import IPNetwork, NetboxAsset, QueryResult


def _get_smallest_subnet(data: QueryResult, ip_network: IPNetwork) -> NetboxAsset | None:
    best_prefix_match: NetboxAsset | None = None
    for prefix in data.prefixes:
        if (
            isinstance(ip_network, IPv4Network)
            and isinstance(prefix.cidr.exploded, IPv4Network)
            and ip_network.subnet_of(prefix.cidr.exploded)
            and (
                best_prefix_match is None
                or (
                    isinstance(best_prefix_match.cidr.exploded, IPv4Network)
                    and prefix.cidr.exploded.subnet_of(best_prefix_match.cidr.exploded)
                )
            )
        ) or (
            isinstance(ip_network, IPv6Network)
            and isinstance(prefix.cidr.exploded, IPv6Network)
            and ip_network.subnet_of(prefix.cidr.exploded)
            and (
                best_prefix_match is None
                or (
                    isinstance(best_prefix_match.cidr.exploded, IPv6Network)
                    and prefix.cidr.exploded.subnet_of(best_prefix_match.cidr.exploded)
                )
            )
        ):
            best_prefix_match = prefix
    return best_prefix_match


def get_subnet_info_for_ip(ip: str, data: QueryResult | None) -> NetboxAsset | None:
    if not data or not data.prefixes:
        return None
    ip_network: IPNetwork = IPv6Network(ip) if ":" in ip else IPv4Network(ip)
    if (subnet := _get_smallest_subnet(data, ip_network)) is not None or not data.ip_addresses:
        return subnet
    for ip_address in data.ip_addresses:
        if ip_address.cidr.exploded == ip_network:
            return ip_address
    return None
