"""轻量 DNS 观察器，用于把 TUN 里的目标 IP 关联回解析域名。

TUN 送到 SOCKS5 的请求通常只带目标 IP。这里观察经 UoT 或直连的
UDP DNS 回包，保存短期的 ``IP -> 域名`` 映射供分流和日志使用。
"""

from __future__ import annotations

import ipaddress
import struct
import threading
import time
from typing import Callable


def _read_name(packet: bytes, offset: int) -> tuple[str, int]:
    labels: list[str] = []
    cursor = offset
    next_offset = None
    jumps = 0
    while True:
        if cursor >= len(packet):
            raise ValueError("DNS 名称超出数据包")
        length = packet[cursor]
        if length == 0:
            cursor += 1
            if next_offset is None:
                next_offset = cursor
            break
        if length & 0xC0 == 0xC0:
            if cursor + 1 >= len(packet):
                raise ValueError("DNS 压缩指针不完整")
            pointer = ((length & 0x3F) << 8) | packet[cursor + 1]
            if next_offset is None:
                next_offset = cursor + 2
            cursor = pointer
            jumps += 1
            if jumps > 16:
                raise ValueError("DNS 名称压缩指针循环")
            continue
        if length > 63 or cursor + 1 + length > len(packet):
            raise ValueError("DNS 标签不完整")
        labels.append(packet[cursor + 1:cursor + 1 + length].decode("idna"))
        cursor += 1 + length
    return ".".join(labels).rstrip(".").lower(), int(next_offset)


def _question(packet: bytes) -> tuple[int, str, int]:
    if len(packet) < 12:
        raise ValueError("DNS 头不完整")
    txid, flags, questions, _answers, _authority, _additional = struct.unpack(
        "!HHHHHH", packet[:12]
    )
    if flags & 0x8000 or questions < 1:
        raise ValueError("不是 DNS 查询")
    name, offset = _read_name(packet, 12)
    if offset + 4 > len(packet):
        raise ValueError("DNS 问题段不完整")
    return txid, name, offset + 4


class DnsObserver:
    def __init__(self, log: Callable[[str], None] | None = None,
                 max_ttl: int = 300):
        self.log = log or (lambda _message: None)
        self.max_ttl = max(1, int(max_ttl))
        self._lock = threading.RLock()
        self._domain_to_ips: dict[str, tuple[float, tuple[str, ...]]] = {}
        # 一个 CDN IP 可能同时服务多个域名；不能把最后一次 DNS 回包当成唯一身份。
        self._ip_to_domains: dict[str, tuple[float, set[str]]] = {}

    def observe_query(self, packet: bytes) -> str | None:
        try:
            _txid, domain, _ = _question(packet)
            return domain or None
        except (ValueError, UnicodeError):
            return None

    def observe_response(self, packet: bytes) -> tuple[str, tuple[str, ...]] | None:
        try:
            if len(packet) < 12:
                raise ValueError("DNS 头不完整")
            _txid, flags, questions, answers, _authority, _additional = struct.unpack(
                "!HHHHHH", packet[:12]
            )
            if not (flags & 0x8000) or questions < 1:
                raise ValueError("不是 DNS 回包")
            domain, offset = _read_name(packet, 12)
            if offset + 4 > len(packet):
                raise ValueError("DNS 问题段不完整")
            offset += 4
            ips: list[str] = []
            ttl = self.max_ttl
            for _ in range(answers):
                _owner, offset = _read_name(packet, offset)
                if offset + 10 > len(packet):
                    raise ValueError("DNS 答案头不完整")
                rtype, rclass, record_ttl, size = struct.unpack(
                    "!HHIH", packet[offset:offset + 10]
                )
                offset += 10
                if offset + size > len(packet):
                    raise ValueError("DNS 答案数据不完整")
                raw = packet[offset:offset + size]
                offset += size
                if rclass != 1 or rtype not in (1, 28):
                    continue
                if (rtype == 1 and size != 4) or (rtype == 28 and size != 16):
                    continue
                ips.append(str(ipaddress.ip_address(raw)))
                ttl = min(ttl, max(1, int(record_ttl)))
            if not domain or not ips:
                return None
            expires = time.monotonic() + min(ttl, self.max_ttl)
            unique = tuple(dict.fromkeys(ips))
            with self._lock:
                self._domain_to_ips[domain] = (expires, unique)
                for ip in unique:
                    previous = self._ip_to_domains.get(ip)
                    domains = set(previous[1]) if previous and previous[0] > time.monotonic() else set()
                    domains.add(domain)
                    self._ip_to_domains[ip] = (expires, domains)
            return domain, unique
        except (ValueError, UnicodeError, OSError):
            return None

    def lookup_ip(self, ip: str) -> str | None:
        try:
            normalized = str(ipaddress.ip_address(ip))
        except ValueError:
            return None
        now = time.monotonic()
        with self._lock:
            item = self._ip_to_domains.get(normalized)
            if not item:
                return None
            if item[0] <= now:
                self._ip_to_domains.pop(normalized, None)
                return None
            return next(iter(item[1])) if len(item[1]) == 1 else None

    def lookup_domain(self, domain: str) -> tuple[str, ...]:
        key = str(domain).rstrip(".").lower()
        now = time.monotonic()
        with self._lock:
            item = self._domain_to_ips.get(key)
            if not item:
                return ()
            if item[0] <= now:
                self._domain_to_ips.pop(key, None)
                return ()
            return item[1]