#!/usr/bin/env python3
"""Fetch RDAP details for all IPv4 and IPv6 prefixes of a country list.

Example:
    python3 fetch_rdap_prefixes.py --country my
"""

from __future__ import annotations

import argparse
from collections import Counter
import ipaddress
import json
import sys
from typing import Any
from urllib.error import HTTPError, URLError
from urllib.request import urlopen

GITHUB_RAW_TEMPLATE = (
    "https://raw.githubusercontent.com/ipverse/country-ip-blocks/"
    "master/country/{country}/aggregated.json"
)
RDAP_TEMPLATE = "http://rdap.apnic.net/ip/{address}"


def fetch_json(url: str, timeout: int = 20) -> dict[str, Any]:
    """Download JSON from a URL and parse it."""
    try:
        with urlopen(url, timeout=timeout) as response:
            body = response.read().decode("utf-8")
        return json.loads(body)
    except HTTPError as exc:
        raise RuntimeError(f"HTTP error {exc.code} for URL: {url}") from exc
    except URLError as exc:
        raise RuntimeError(f"Network error for URL {url}: {exc.reason}") from exc
    except json.JSONDecodeError as exc:
        raise RuntimeError(f"Invalid JSON response from URL: {url}") from exc


def all_prefixes(country: str, count: int | None = None) -> dict[str, list[str]]:
    url = GITHUB_RAW_TEMPLATE.format(country=country.lower())
    data = fetch_json(url)

    prefixes = data.get("prefixes", {})
    ipv4 = prefixes.get("ipv4")
    ipv6 = prefixes.get("ipv6")
    if not isinstance(ipv4, list):
        raise RuntimeError("Missing or invalid prefixes.ipv4 array in country JSON")
    if not isinstance(ipv6, list):
        raise RuntimeError("Missing or invalid prefixes.ipv6 array in country JSON")

    ipv4_list = [str(item) for item in ipv4]
    ipv6_list = [str(item) for item in ipv6]
    if count is not None:
        ipv4_list = ipv4_list[:count]
        ipv6_list = ipv6_list[:count]

    return {
        "ipv4": ipv4_list,
        "ipv6": ipv6_list,
    }


def rdap_for_prefix(prefix: str) -> dict[str, str]:
    query_ip = prefix.split("/", 1)[0]
    rdap_url = RDAP_TEMPLATE.format(address=query_ip)
    rdap = fetch_json(rdap_url)

    return {
        "prefix": prefix,
        "query_ip": query_ip,
        "handle": str(rdap.get("handle", "")),
        "name": str(rdap.get("name", "")),
        "country": str(rdap.get("country", "")),
        "type": str(rdap.get("type", "")),
        "startAddress": str(rdap.get("startAddress", "")),
        "endAddress": str(rdap.get("endAddress", "")),
    }


def prefix_address_count(prefix: str) -> int:
    """Return the total number of addresses represented by a CIDR prefix."""
    network = ipaddress.ip_network(prefix, strict=False)
    return network.num_addresses


def prefix_site56_count(prefix: str) -> int:
    """Return how many full /56 IPv6 site prefixes fit in this prefix."""
    network = ipaddress.ip_network(prefix, strict=False)
    if network.version != 6:
        return 0
    if network.prefixlen > 56:
        return 0
    return 1 << (56 - network.prefixlen)


def main() -> int:
    parser = argparse.ArgumentParser(
        description=(
            "Fetch RDAP info for all IPv4 and IPv6 prefixes from ipverse country list"
        )
    )
    parser.add_argument("--country", default="my", help="Country code (default: my)")
    parser.add_argument(
        "--count",
        type=int,
        default=None,
        help="Optional limit: process first N prefixes for both IPv4 and IPv6",
    )
    args = parser.parse_args()

    if args.count is not None and args.count <= 0:
        print("Error: --count must be greater than 0", file=sys.stderr)
        return 2

    try:
        prefixes_by_family = all_prefixes(args.country, args.count)
    except RuntimeError as exc:
        print(f"Error loading prefixes: {exc}", file=sys.stderr)
        return 1

    print("ip_version|prefix|query_ip|handle|name|country|type|startAddress|endAddress")

    name_counts_ipv4: Counter[str] = Counter()
    name_counts_ipv6: Counter[str] = Counter()
    name_counts_all: Counter[str] = Counter()
    name_ipv4_addresses: Counter[str] = Counter()
    name_ipv6_site56: Counter[str] = Counter()

    for ip_version, prefixes in (("ipv4", prefixes_by_family["ipv4"]), ("ipv6", prefixes_by_family["ipv6"])):
        for prefix in prefixes:
            try:
                row = rdap_for_prefix(prefix)
            except RuntimeError as exc:
                # Keep processing remaining prefixes if one query fails.
                print(f"{ip_version}|{prefix}|{prefix.split('/', 1)[0]}|ERROR|{exc}||||")
                continue

            print(
                "|".join(
                    [
                        ip_version,
                        row["prefix"],
                        row["query_ip"],
                        row["handle"],
                        row["name"],
                        row["country"],
                        row["type"],
                        row["startAddress"],
                        row["endAddress"],
                    ]
                )
            )

            normalized_name = row["name"].strip() or "(unknown)"
            address_count = prefix_address_count(prefix)
            if ip_version == "ipv4":
                name_counts_ipv4[normalized_name] += 1
                name_ipv4_addresses[normalized_name] += address_count
            else:
                name_counts_ipv6[normalized_name] += 1
                name_ipv6_site56[normalized_name] += prefix_site56_count(prefix)
            name_counts_all[normalized_name] += 1

    print()
    print("summary_metric|value")
    print(f"unique_names_ipv4|{len(name_counts_ipv4)}")
    print(f"unique_names_ipv6|{len(name_counts_ipv6)}")
    print(f"unique_names_total|{len(name_counts_all)}")
    print(f"total_prefixes_ipv4|{sum(name_counts_ipv4.values())}")
    print(f"total_prefixes_ipv6|{sum(name_counts_ipv6.values())}")
    print(f"total_prefixes_all|{sum(name_counts_all.values())}")
    print(f"total_ipv6_site56|{sum(name_ipv6_site56.values())}")

    print()
    print("name|prefix_count_total|ipv4_addresses|ipv6_site56")
    for name, prefix_count in sorted(
        name_counts_all.items(), key=lambda item: (-item[1], item[0].lower())
    ):
        print(
            f"{name}|{prefix_count}|{name_ipv4_addresses[name]}|{name_ipv6_site56[name]}"
        )

    return 0


if __name__ == "__main__":
    raise SystemExit(main())