"""
Advertise the dashboard on the LAN under its own name, e.g. `agibot.local`.

Why this exists
---------------
The link people type should be about the robot, not about whatever the PC
happens to be called. Renaming the computer would do it, but that needs
administrator rights and a reboot and changes the machine's identity for
everything else. Publishing an extra name over mDNS is additive, reversible and
needs neither: the computer keeps its own name, and the advertised name exists
only while the dashboard is running.

Why stdlib and not the `zeroconf` package
-----------------------------------------
Two measured constraints on this machine, either of which would have made a
separate helper process silently useless:

  * Windows' Public firewall profile is BlockInbound and every inbound rule for
    UDP 5353 is *program*-scoped, not port-scoped. The built-in "mDNS (UDP-In)"
    rule is scoped to svchost.exe/dnscache and grants a Python process nothing.
    The only interpreter with a blanket inbound allow is the base Python that
    already serves the dashboard - so the responder has to live *inside* that
    process to be reachable from another device.
  * A same-host HTTP fetch to the PC's own LAN address is routed via loopback
    and skips inbound firewall filtering entirely, so a helper running under the
    wrong interpreter passes every local test while remaining invisible to a
    phone. Running in-process removes that whole failure mode.

Being stdlib-only also means no extra dependency to install in the interpreter
that matters.

Scope
-----
Only the `<name>.local` form is published. A bare single-label `http://agibot:`
appears to work on Windows because Windows appends `.local` to single-label
names, but it is unreliable (its negative DNS cache can poison the name for
minutes) and no iPhone, Android, Mac or Linux client resolves it. Publishing it
would be advertising something that only works on one machine.
"""

from __future__ import annotations

import logging
import socket
import struct
import threading
import time

log = logging.getLogger("x2.announce")

MDNS_ADDR = "224.0.0.251"
MDNS_PORT = 5353
TTL = 120                    # seconds a client may cache the record
QTYPE_A = 1
QTYPE_ANY = 255
CLASS_IN = 1
CACHE_FLUSH = 0x8000         # tells clients this is the authoritative answer
UNICAST_RESPONSE = 0x8000    # the "QU" bit on a question's class


def _encode_name(name: str) -> bytes:
    out = bytearray()
    for label in name.rstrip(".").split("."):
        raw = label.encode("utf-8")
        out.append(len(raw))
        out += raw
    out.append(0)
    return bytes(out)


def _decode_name(data: bytes, offset: int) -> tuple[str, int]:
    """Read a DNS name, following compression pointers. Returns (name, next_offset)."""
    labels: list[str] = []
    jumped = False
    end = offset
    hops = 0

    while True:
        if offset >= len(data) or hops > 20:
            break
        length = data[offset]

        if length & 0xC0 == 0xC0:                      # compression pointer
            if offset + 1 >= len(data):
                break
            pointer = struct.unpack_from("!H", data, offset)[0] & 0x3FFF
            if not jumped:
                end = offset + 2
            offset = pointer
            jumped = True
            hops += 1
            continue

        offset += 1
        if length == 0:
            if not jumped:
                end = offset
            break
        labels.append(data[offset:offset + length].decode("utf-8", "replace"))
        offset += length

    return ".".join(labels), end


def _build_response(name: str, ip: str) -> bytes:
    header = struct.pack(
        "!HHHHHH",
        0,          # ID - always 0 in mDNS
        0x8400,     # QR=1 (response), AA=1 (authoritative)
        0,          # no questions echoed
        1,          # one answer
        0, 0,
    )
    answer = (
        _encode_name(name)
        + struct.pack("!HHIH", QTYPE_A, CLASS_IN | CACHE_FLUSH, TTL, 4)
        + socket.inet_aton(ip)
    )
    return header + answer


def _build_query(name: str) -> bytes:
    header = struct.pack("!HHHHHH", 0, 0x0000, 1, 0, 0, 0)
    return header + _encode_name(name) + struct.pack("!HH", QTYPE_A, CLASS_IN)


class Announcer:
    """Answers mDNS A queries for one name, for as long as it runs."""

    def __init__(self, name: str, port: int, address_provider):
        self.name = f"{name.strip().strip('.').lower()}.local"
        self.port = port
        self._address_of = address_provider
        self._sock: socket.socket | None = None
        self._thread: threading.Thread | None = None
        self._stop = threading.Event()
        self.ip: str | None = None
        self.answered = 0
        self.error = ""
        self.conflict = False

    # -- lifecycle ----------------------------------------------------------

    def start(self) -> bool:
        self.ip = self._address_of()
        if not self.ip:
            self.error = "No LAN address to advertise"
            log.warning("mDNS: %s", self.error)
            return False

        try:
            self._sock = self._open_socket()
        except OSError as exc:
            self.error = f"Could not open UDP {MDNS_PORT}: {exc}"
            log.warning("mDNS: %s", self.error)
            return False

        # RFC 6762 probing: if something else already owns the name, do not
        # fight it - two hosts answering for one name gives whichever reply
        # arrives first, which is worse than not advertising at all.
        if self._name_taken():
            self.conflict = True
            self.error = f"{self.name} is already claimed by another device on this network"
            log.warning("mDNS: %s", self.error)
            self._close()
            return False

        self._stop.clear()
        self._thread = threading.Thread(target=self._serve, name="mdns-announce", daemon=True)
        self._thread.start()
        log.info("mDNS: advertising %s -> %s:%s", self.name, self.ip, self.port)
        return True

    def stop(self) -> None:
        self._stop.set()
        self._close()
        if self._thread:
            self._thread.join(timeout=3.0)
            self._thread = None

    def _close(self) -> None:
        if self._sock is not None:
            try:
                self._sock.close()
            except OSError:
                pass
            self._sock = None

    # -- socket -------------------------------------------------------------

    def _open_socket(self) -> socket.socket:
        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        # Several responders normally share 5353 (Windows' own, Bonjour, and
        # anything Chrome is doing); SO_REUSEPORT where available lets us join
        # them rather than fail.
        if hasattr(socket, "SO_REUSEPORT"):
            try:
                sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1)
            except OSError:
                pass
        sock.bind(("", MDNS_PORT))

        membership = socket.inet_aton(MDNS_ADDR) + socket.inet_aton(self.ip)
        sock.setsockopt(socket.IPPROTO_IP, socket.IP_ADD_MEMBERSHIP, membership)
        sock.setsockopt(socket.IPPROTO_IP, socket.IP_MULTICAST_TTL, 255)
        sock.setsockopt(socket.IPPROTO_IP, socket.IP_MULTICAST_IF, socket.inet_aton(self.ip))
        sock.settimeout(1.0)
        return sock

    def _name_taken(self) -> bool:
        """Ask whether anyone else already answers for our name."""
        probe = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        try:
            probe.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            probe.setsockopt(socket.IPPROTO_IP, socket.IP_MULTICAST_TTL, 255)
            probe.settimeout(0.4)
            probe.sendto(_build_query(self.name), (MDNS_ADDR, MDNS_PORT))

            deadline = time.time() + 0.9
            while time.time() < deadline:
                try:
                    data, addr = probe.recvfrom(4096)
                except socket.timeout:
                    continue
                except OSError:
                    break
                if addr[0] == self.ip:
                    continue                       # our own machine, ignore
                if self._is_answer_for_us(data):
                    log.warning("mDNS: %s already answered by %s", self.name, addr[0])
                    return True
        except OSError:
            pass
        finally:
            probe.close()
        return False

    def _is_answer_for_us(self, data: bytes) -> bool:
        try:
            _, flags, qd, an, _, _ = struct.unpack_from("!HHHHHH", data, 0)
        except struct.error:
            return False
        if not (flags & 0x8000) or an == 0:
            return False

        offset = 12
        for _ in range(qd):
            _, offset = _decode_name(data, offset)
            offset += 4
        for _ in range(an):
            name, offset = _decode_name(data, offset)
            try:
                rtype, _, _, rdlen = struct.unpack_from("!HHIH", data, offset)
            except struct.error:
                return False
            offset += 10 + rdlen
            if rtype == QTYPE_A and name.lower() == self.name:
                return True
        return False

    # -- serving ------------------------------------------------------------

    def _serve(self) -> None:
        last_ip_check = time.time()

        while not self._stop.is_set():
            sock = self._sock
            if sock is None:
                return
            try:
                data, addr = sock.recvfrom(4096)
            except socket.timeout:
                data = None
            except OSError:
                if not self._stop.is_set():
                    log.debug("mDNS: socket closed")
                return

            # DHCP leases change; an advertiser pinned to a stale address will
            # confidently point the name at nothing.
            if time.time() - last_ip_check > 30:
                last_ip_check = time.time()
                current = self._address_of()
                if current and current != self.ip:
                    log.info("mDNS: address changed %s -> %s, re-advertising", self.ip, current)
                    self.ip = current

            if not data:
                continue

            try:
                self._handle(data, addr)
            except Exception as exc:            # never let one packet stop us
                log.debug("mDNS: bad packet from %s: %s", addr, exc)

    def _handle(self, data: bytes, addr) -> None:
        try:
            _, flags, qdcount, _, _, _ = struct.unpack_from("!HHHHHH", data, 0)
        except struct.error:
            return
        if flags & 0x8000 or qdcount == 0:      # a response, not a question
            return

        offset = 12
        for _ in range(qdcount):
            name, offset = _decode_name(data, offset)
            try:
                qtype, qclass = struct.unpack_from("!HH", data, offset)
            except struct.error:
                return
            offset += 4

            if name.lower() != self.name:
                continue
            if qtype not in (QTYPE_A, QTYPE_ANY):
                continue

            response = _build_response(self.name, self.ip)
            wants_unicast = bool(qclass & UNICAST_RESPONSE)
            sock = self._sock
            if sock is None:
                return
            try:
                if wants_unicast:
                    sock.sendto(response, addr)
                else:
                    # Multicast so every listener refreshes its cache, and
                    # unicast too because some stacks only accept the direct
                    # reply.
                    sock.sendto(response, (MDNS_ADDR, MDNS_PORT))
                    sock.sendto(response, addr)
                self.answered += 1
            except OSError:
                pass

    # -- introspection ------------------------------------------------------

    def status(self) -> dict:
        return {
            "name": self.name,
            "url": f"http://{self.name}:{self.port}",
            "ip": self.ip,
            "running": bool(self._thread and self._thread.is_alive()),
            "answered": self.answered,
            "conflict": self.conflict,
            "error": self.error,
        }