"""DNS resolution check — pure probe returning a result dict.

With no resolver, the name is looked up the way the container looks names
up (getaddrinfo: /etc/hosts, then the resolvers in resolv.conf). With one,
an A and an AAAA query go straight to that server over UDP, and the time
it took to answer is the check's response time — so "does my local DNS
still answer for public names" and "does 1.1.1.1" are separate monitors.
"""
from __future__ import annotations
import asyncio
import ipaddress
import logging
import os
import socket
import struct
import time

logger = logging.getLogger(__name__)

QUERY_TIMEOUT_S = 5.0
_RCODES = {1: "FORMERR", 2: "SERVFAIL", 3: "NXDOMAIN", 4: "NOTIMP", 5: "REFUSED"}
_A, _AAAA = 1, 28


async def dns_check(address: str, expected_ip: str | None = None,
                    resolver: str | None = None) -> dict:
    """Resolve `address`. Returns {is_up, resolved_ip, error_msg} and, when a
    resolver is given, response_ms.

    Up = resolution succeeded AND (no expected_ip, or at least one returned
    record matches expected_ip). resolved_ip is the first A/AAAA record.
    """
    if resolver:
        result = await _query_server(address, resolver)
    else:
        result = await _resolve(address)
    out = {"response_ms": result["ms"]} if result.get("ms") is not None else {}
    if not result["resolved"]:
        return {**out, "is_up": False, "resolved_ip": None,
                "error_msg": result.get("error") or "Resolution failed"}

    resolved_ip = result["ip"]
    if expected_ip and expected_ip not in result["all_ips"]:
        return {
            **out,
            "is_up": False,
            "resolved_ip": resolved_ip,
            "error_msg": f"Expected {expected_ip} not in {', '.join(result['all_ips'])}",
        }
    return {**out, "is_up": True, "resolved_ip": resolved_ip, "error_msg": None}


async def _resolve(address: str) -> dict:
    """Resolve hostname. Returns all A/AAAA records; first is the resolved_ip."""
    try:
        loop = asyncio.get_running_loop()
        infos = await loop.run_in_executor(
            None, socket.getaddrinfo, address, None, socket.AF_UNSPEC, socket.SOCK_STREAM
        )
        if infos:
            all_ips = [info[4][0] for info in infos]
            return {"resolved": True, "ip": all_ips[0], "all_ips": all_ips}
        return {"resolved": False, "ip": None, "all_ips": []}
    except Exception as exc:
        logger.debug("DNS resolution failed for %s: %s", address, exc)
        return {"resolved": False, "ip": None, "all_ips": []}


async def _query_server(address: str, resolver: str) -> dict:
    """Ask `resolver` directly for A and AAAA, together; the slower answer is the time."""
    started = time.monotonic()
    answers = await asyncio.gather(_ask(address, resolver, _A), _ask(address, resolver, _AAAA),
                                   return_exceptions=True)
    ms = round((time.monotonic() - started) * 1000, 1)
    ips, errors = [], []
    for got in answers:
        (errors if isinstance(got, Exception) else ips).append(got)
    all_ips = [ip for batch in ips for ip in batch]
    if all_ips:
        return {"resolved": True, "ip": all_ips[0], "all_ips": all_ips, "ms": ms}
    if errors:
        # A timeout on both queries leaves no time worth recording.
        timed_out = all(isinstance(e, TimeoutError) for e in errors) and not ips
        return {"resolved": False, "ip": None, "all_ips": [], "ms": None if timed_out else ms,
                "error": f"{resolver}: {errors[0]}"}
    return {"resolved": False, "ip": None, "all_ips": [], "ms": ms,
            "error": f"{resolver} has no A or AAAA record for {address}"}


class _Reply(asyncio.DatagramProtocol):
    def __init__(self, qid: int):
        self.qid = qid
        self.done = asyncio.get_running_loop().create_future()

    def datagram_received(self, data, addr):
        if len(data) >= 2 and struct.unpack("!H", data[:2])[0] == self.qid and not self.done.done():
            self.done.set_result(data)

    def error_received(self, exc):
        if not self.done.done():
            self.done.set_exception(exc)


async def _ask(name: str, resolver: str, qtype: int) -> list[str]:
    """One UDP query; the addresses of `qtype` in the answer, or an exception saying why not."""
    qid = struct.unpack("!H", os.urandom(2))[0]
    loop = asyncio.get_running_loop()
    transport, proto = await loop.create_datagram_endpoint(
        lambda: _Reply(qid), remote_addr=(resolver, 53))
    try:
        transport.sendto(build_query(qid, name, qtype))
        try:
            data = await asyncio.wait_for(proto.done, QUERY_TIMEOUT_S)
        except asyncio.TimeoutError:
            raise TimeoutError(f"no answer within {QUERY_TIMEOUT_S:.0f}s") from None
    finally:
        transport.close()
    return parse_answer(data, qtype)


def build_query(qid: int, name: str, qtype: int) -> bytes:
    """A standard recursive query for `name`."""
    header = struct.pack("!HHHHHH", qid, 0x0100, 1, 0, 0, 0)   # RD set, one question
    labels = b"".join(bytes([len(p)]) + p.encode("idna") for p in name.strip(".").split(".") if p)
    return header + labels + b"\x00" + struct.pack("!HH", qtype, 1)


def parse_answer(data: bytes, qtype: int) -> list[str]:
    """The `qtype` addresses in a reply; ValueError naming the server's error code."""
    if len(data) < 12:
        raise ValueError("short reply")
    _, flags, qd, an, _, _ = struct.unpack("!HHHHHH", data[:12])
    rcode = flags & 0x000F
    if rcode:
        raise ValueError(_RCODES.get(rcode, f"error code {rcode}"))
    at = 12
    for _ in range(qd):
        at = _skip_name(data, at) + 4
    ips = []
    for _ in range(an):
        at = _skip_name(data, at)
        rtype, _, _, length = struct.unpack("!HHIH", data[at:at + 10])
        at += 10
        rdata = data[at:at + length]
        at += length
        if rtype == qtype == _A and length == 4:
            ips.append(str(ipaddress.IPv4Address(rdata)))
        elif rtype == qtype == _AAAA and length == 16:
            ips.append(str(ipaddress.IPv6Address(rdata)))
    return ips


def _skip_name(data: bytes, at: int) -> int:
    while True:
        n = data[at]
        if n == 0:
            return at + 1
        if n & 0xC0 == 0xC0:   # compression pointer ends the name
            return at + 2
        at += n + 1