"""Merge and deduplicate IP address ranges from input files."""

# PYTHON_ARGCOMPLETE_OK

import argparse
import fileinput

import argcomplete
from netaddr import IPNetwork, IPSet


def digit_str_zfill(digit_str: str, group: int) -> str:
    """Add leading zeros to groups of digits.

    Args:
        digit_str: String of digits to pad
        group: Size of digit groups

    Returns:
        Zero-padded digit string
    """
    zfill = (len(digit_str) + group - 1) // group * group
    return " ".join(
        digit_str.zfill(zfill)[i : i + group] for i in range(0, len(digit_str), group)
    )


def digit_to_binary(digit_str: str, base: int = 16, group: int = 4) -> str:
    """Convert digits to binary representation.

    Args:
        digit_str: String of digits to convert
        base: Base of input digits
        group: Size of output binary groups

    Returns:
        Binary string representation
    """
    try:
        binary_str = bin(int(digit_str, base)).removeprefix("0b")
        return digit_str_zfill(binary_str, group)
    except ValueError:
        return digit_str


def ip_network_to_binary(ip_network: IPNetwork) -> str:
    """Convert IP network to binary representation.

    Args:
        ip_network: IP network to convert

    Returns:
        Binary string representation
    """
    sep, base, group = (":", 16, 4) if ip_network.version == 6 else (".", 10, 8)
    binary_addr = sep.join(
        digit_to_binary(digit_str, base, group)
        for digit_str in str(ip_network.network).split(sep)
    )
    return f"{binary_addr}/{ip_network.prefixlen}"


def ip_network_zfill(ip_network: IPNetwork) -> str:
    sep, group = (":", 4) if ip_network.version == 6 else (".", 3)
    zfill_addr = sep.join(
        digit_str_zfill(digit_str, group)
        for digit_str in str(ip_network.network).split(sep)
    )
    return f"{zfill_addr}/{ip_network.prefixlen}"


def merge_ip_ranges(ip_range_files: list[str]) -> tuple[IPSet, IPSet]:
    """Merge IP ranges from input files.

    Args:
        ip_range_files: List of files containing IP ranges

    Returns:
        Tuple of (IPv4Set, IPv6Set) containing merged ranges
    """
    ipv4_set, ipv6_set = IPSet(), IPSet()
    try:
        # Read from input sources (either files or stdin)
        for ip_range in fileinput.input(files=ip_range_files):
            ip_range = ip_range.strip()
            if not ip_range:
                continue
            ip_network = IPNetwork(ip_range)
            if ip_network.version == 4:
                ipv4_set.add(ip_network)
            elif ip_network.version == 6:
                ipv6_set.add(ip_network)
    except KeyboardInterrupt:
        print()

    return ipv4_set, ipv6_set


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="merge IP ranges from files or standard input."
    )
    parser.add_argument(
        "-4",
        "--ipv4",
        action="store_true",
        default=False,
        help="output only IPv4 addresses",
    )
    parser.add_argument(
        "-6",
        "--ipv6",
        action="store_true",
        default=False,
        help="output only IPv6 addresses",
    )
    parser.add_argument(
        "-b",
        "--binary",
        action="store_true",
        default=False,
        help="output addresses in binary format.",
    )
    parser.add_argument(
        "-z",
        "--zfill",
        action="store_true",
        default=False,
        help="output addresses prefixed with zero.",
    )
    parser.add_argument(
        "files",
        nargs="*",
        metavar="FILE",
        help="input files containing IP ranges, if not specified, reads from stdin.",
    )
    argcomplete.autocomplete(parser)

    return parser.parse_args()


def main() -> None:
    """Parse arguments and print merged/deduplicated IP ranges."""
    args = parse_args()

    ipv4_set, ipv6_set = merge_ip_ranges(args.files)
    if args.ipv4:
        merged_ranges = ipv4_set
    elif args.ipv6:
        merged_ranges = ipv6_set
    else:
        merged_ranges = ipv4_set.union(ipv6_set)

    for cidr in merged_ranges.iter_cidrs():
        if args.binary:
            print(ip_network_to_binary(cidr))
        elif args.zfill:
            print(ip_network_zfill(cidr))
        else:
            print(str(cidr))


if __name__ == "__main__":
    main()
