#!/usr/bin/env python3
"""WHOIS subnet querying and IP consolidation."""

from __future__ import annotations

import ipaddress
import json
import os
import re
from typing import Iterable

from ip_utils import ip_version, parse_ip

# WHOIS Cache file path
_WHOIS_CACHE_PATH = os.path.join(
    os.path.dirname(os.path.abspath(__file__)), "data", "whois_cache.json"
)

_whois_cache: dict[str, str] = {}


def load_whois_cache() -> None:
    global _whois_cache
    if not _whois_cache and os.path.exists(_WHOIS_CACHE_PATH):
        try:
            with open(_WHOIS_CACHE_PATH, "r", encoding="utf-8") as f:
                _whois_cache = json.load(f)
        except Exception:
            _whois_cache = {}


def save_whois_cache() -> None:
    if _whois_cache:
        os.makedirs(os.path.dirname(_WHOIS_CACHE_PATH), exist_ok=True)
        try:
            with open(_WHOIS_CACHE_PATH, "w", encoding="utf-8") as f:
                json.dump(_whois_cache, f, indent=2)
        except Exception:
            pass


def query_whois_subnet(ip: str) -> str | None:
    """Query whois.radb.net for the route/subnet of an IP address."""
    addr = parse_ip(ip)
    if addr is None:
        return None

    load_whois_cache()
    if ip in _whois_cache:
        return _whois_cache[ip] or None

    import socket
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.settimeout(3.0)
        s.connect(("whois.radb.net", 43))
        s.sendall(f"{ip}\r\n".encode())
        res = ""
        while True:
            chunk = s.recv(4096).decode(errors="ignore")
            if not chunk:
                break
            res += chunk
        s.close()

        routes = re.findall(r"^(?:route|route6):\s*(\S+)",
                            res, re.IGNORECASE | re.MULTILINE)
        if routes:
            routes_parsed = []
            for r in routes:
                try:
                    routes_parsed.append(ipaddress.ip_network(r))
                except ValueError:
                    continue
            if routes_parsed:
                best_route = max(routes_parsed, key=lambda n: n.prefixlen)
                subnet_str = str(best_route)
                _whois_cache[ip] = subnet_str
                save_whois_cache()
                return subnet_str

    except Exception:
        pass

    subnet_str = f"{ip}/32" if addr.version == 4 else f"{ip}/128"
    _whois_cache[ip] = subnet_str
    save_whois_cache()
    return subnet_str


def consolidate_by_subnet(ips: Iterable[str], threshold: int = 3) -> set[str]:
    """Group IPs by WHOIS subnet; if >= threshold IPs in a subnet, block the subnet.

    Otherwise, block individual IPs.
    """
    ip_list = list(ips)
    subnet_to_ips: dict[str, set[str]] = {}
    ip_to_subnet: dict[str, str] = {}

    for ip in ip_list:
        subnet = query_whois_subnet(ip)
        if subnet:
            try:
                net = ipaddress.ip_network(subnet, strict=False)
                addr = ipaddress.ip_address(ip)
                if addr in net:
                    subnet_to_ips.setdefault(subnet, set()).add(ip)
                    ip_to_subnet[ip] = subnet
                    continue
            except ValueError:
                pass
        fallback_net = f"{ip}/32" if ip_version(ip) == 4 else f"{ip}/128"
        subnet_to_ips.setdefault(fallback_net, set()).add(ip)
        ip_to_subnet[ip] = fallback_net

    blocked_targets: set[str] = set()
    for subnet, mapped_ips in subnet_to_ips.items():
        if len(mapped_ips) >= threshold:
            blocked_targets.add(subnet)
        else:
            for ip in mapped_ips:
                blocked_targets.add(
                    f"{ip}/32" if ip_version(ip) == 4 else f"{ip}/128")

    return blocked_targets


# Re-export path for backward compat
WHOIS_CACHE_PATH = _WHOIS_CACHE_PATH