"""A pure python implementation of multicast DNS service discovery.

Licensed under LGPL-2.1-or-later; see COPYING for details. This file is
part of a continuously modified work; modification dates are recorded
in the project's git history.
"""

from __future__ import annotations

import asyncio
import itertools
import socket
import threading
from typing import TYPE_CHECKING, cast

from ._logger import log
from ._record_update import RecordUpdate
from ._utils.asyncio import get_running_loop, run_coro_with_timeout
from ._utils.net import (
    InterfacesType,
    IPVersion,
    add_interface,
    add_multicast_member,
    drop_multicast_member,
    new_listen_socket,
    normalize_interface_choice,
)
from ._utils.time import current_time_millis
from .const import _CACHE_CLEANUP_INTERVAL

if TYPE_CHECKING:
    from ._core import Zeroconf


from ._listener import AsyncListener
from ._transport import _strip_zone, _WrappedTransport, make_wrapped_transport

_CLOSE_TIMEOUT = 3000  # ms


def _interface_key(interface: str | tuple[tuple[str, int, int], int]) -> tuple[str, int]:
    """Return the (address, scope_id) an interface choice maps to, for diffing.

    Must produce the same key shape as ``_WrappedTransport.interface_key`` so
    the desired set (from ``normalize_interface_choice``) and the current set
    (from the bound senders) diff against each other.
    """
    if isinstance(interface, tuple):
        return (_strip_zone(interface[0][0]), interface[0][2])
    return (interface, 0)


def _listen_socket_supports(
    listen_socket: socket.socket, interface: str | tuple[tuple[str, int, int], int]
) -> bool:
    """Whether the fixed-family listen socket can join this interface's group."""
    if isinstance(interface, tuple):
        # An IPv6 interface can only be joined on an AF_INET6 socket.
        return listen_socket.family == socket.AF_INET6
    if listen_socket.family != socket.AF_INET6:
        # An IPv4 interface on an AF_INET socket.
        return True
    # An IPv4 interface on an AF_INET6 socket: only when it is dual-stack.
    supported = True
    try:
        supported = not listen_socket.getsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY)
    except OSError as exc:
        # Reading IPV6_V6ONLY essentially never fails on a valid AF_INET6
        # socket. Assume dual-stack rather than abort the rescan; returning
        # False instead could loop rebuilds if the rebuilt socket's read also
        # fails. Log at warning, not debug, because if the socket really were
        # v6-only this skips a needed rebuild and leaves an added IPv4 family
        # unreceivable, which is worth surfacing.
        log.warning("Unable to read IPV6_V6ONLY, assuming dual-stack: %s", exc)
    return supported


def _without_transport(
    wrappers: list[_WrappedTransport], transport: asyncio.DatagramTransport
) -> list[_WrappedTransport]:
    """Return the wrappers whose underlying transport is not ``transport``."""
    return [wrapped for wrapped in wrappers if wrapped.transport is not transport]


class AsyncEngine:
    """An engine wraps sockets in the event loop."""

    __slots__ = (
        "_cleanup_timer",
        "_listen_socket",
        "_listen_transport",
        "_respond_sockets",
        "_setup_task",
        "loop",
        "protocols",
        "readers",
        "running_future",
        "senders",
        "zc",
    )

    def __init__(
        self,
        zeroconf: Zeroconf,
        listen_socket: socket.socket | None,
        respond_sockets: list[socket.socket],
    ) -> None:
        self.loop: asyncio.AbstractEventLoop | None = None
        self.zc = zeroconf
        self.protocols: list[AsyncListener] = []
        self.readers: list[_WrappedTransport] = []
        self.senders: list[_WrappedTransport] = []
        self.running_future: asyncio.Future[bool | None] | None = None
        self._listen_socket = listen_socket
        self._listen_transport: _WrappedTransport | None = None
        self._respond_sockets = respond_sockets
        self._cleanup_timer: asyncio.TimerHandle | None = None
        self._setup_task: asyncio.Task[None] | None = None

    async def async_update_interfaces(
        self,
        interfaces: InterfacesType,
        ip_version: IPVersion,
        apple_p2p: bool,
    ) -> bool:
        """Reconcile sender/reader sockets to the live interface set.

        Adds a per-interface responder socket for each interface that
        appeared and tears down the socket for each interface that
        disappeared, diffing on the bound address. A Default single-family
        instance's dual-use listen/responder socket is converted to a pure
        listener when moving to an explicit set; otherwise the shared listen
        socket is left intact. Returns whether any responder socket was
        added, so the caller can skip re-announcing when nothing appeared.
        """
        assert self.loop is not None
        try:
            normalized = normalize_interface_choice(interfaces, ip_version)
        except RuntimeError as exc:
            # An All/Default instance can transiently resolve to zero addresses
            # during adapter churn, where normalize_interface_choice raises
            # instead of returning an empty set. Treat that as a logged no-op so
            # a momentary down state doesn't crash a caller's adapter-change
            # handler (best-effort contract); the next rescan reconciles.
            log.warning("Skipping interface update; no interfaces available: %s", exc)
            return False
        desired = {_interface_key(interface): interface for interface in normalized}
        current = {wrapped.interface_key: wrapped for wrapped in self.senders}
        listen_transport = self._listen_transport
        listen_socket = listen_transport.sock if listen_transport is not None else None

        # The listen socket's family is fixed at construction, so a desired
        # interface of another family (e.g. an IPv6 interface added to an IPv4
        # instance) needs a fresh listen socket before senders are reconciled,
        # otherwise the current senders would be torn down with no replacements
        # bound.
        needs_rebuild = listen_socket is not None and any(
            not _listen_socket_supports(listen_socket, interface) for interface in desired.values()
        )

        # A Default single-family instance shares the listen socket as its only
        # sender (the dual-use socket). Moving it to an explicit interface set
        # abandons that optimization: rebuild it as a pure listener (its existing
        # group memberships would otherwise collide with the new per-interface
        # joins), which also stops it responding so it can't double-announce on
        # the overlapping interface. Drop it from the diff's view so the desired
        # interfaces are added fresh; the actual sender removal is done by the
        # rebuild once it succeeds (so a failed rebuild leaves senders intact).
        # The no-arg refresh of a Default instance leaves desired == {its
        # interface} and so neither converts nor rebuilds.
        if listen_transport is not None and any(
            wrapped.transport is listen_transport.transport for wrapped in self.senders
        ):
            listen_key = listen_transport.interface_key
            if any(key != listen_key for key in desired):
                current.pop(listen_key, None)
                needs_rebuild = True

        if needs_rebuild:
            try:
                await self._async_rebuild_listen_socket(apple_p2p, desired, current)
            except (OSError, RuntimeError) as exc:
                # A fresh wildcard bind / endpoint creation can transiently fail
                # during adapter churn. The rebuild raises before tearing down
                # the old listen socket, so state is unchanged; log and no-op
                # rather than crash the caller's handler (best-effort contract).
                log.warning("Skipping interface update; listen socket rebuild failed: %s", exc)
                return False
            listen_transport = self._listen_transport
            listen_socket = listen_transport.sock if listen_transport is not None else None

        for bind_address, wrapped in current.items():
            if bind_address in desired:
                continue
            if listen_transport is not None and wrapped.transport is listen_transport.transport:
                # The shared listen / dual-use socket is not a per-interface
                # sender; leaving the group or closing it would break receive.
                continue
            # After a rebuild, listen_socket is the new socket, which this gone
            # interface never joined (its membership died with the old socket);
            # the leave is then a benign no-op that drop_multicast_member swallows.
            self._async_close_sender(wrapped, listen_socket)

        added = False
        for bind_address, interface in desired.items():
            if bind_address in current:
                continue
            if await self._async_add_interface(interface, listen_socket, apple_p2p):
                added = True
        return added

    def close(self) -> None:
        """Close from sync context.

        While it is not expected during normal operation,
        this function may raise EventLoopBlocked if the underlying
        call to `_async_close` cannot be completed.
        """
        assert self.loop is not None
        # Guard against Zeroconf.close() being called from the eventloop
        if get_running_loop() == self.loop:
            self._async_shutdown()
            return
        if not self.loop.is_running():
            return
        run_coro_with_timeout(self._async_close(), self.loop, _CLOSE_TIMEOUT)

    def setup(
        self,
        loop: asyncio.AbstractEventLoop,
        loop_thread_ready: threading.Event | None,
    ) -> None:
        """Set up the instance."""
        self.loop = loop
        self.running_future = loop.create_future()
        self._setup_task = self.loop.create_task(self._async_setup(loop_thread_ready))

    async def _async_add_interface(
        self,
        interface: str | tuple[tuple[str, int, int], int],
        listen_socket: socket.socket | None,
        apple_p2p: bool,
    ) -> bool:
        """Join the multicast group and adopt a responder socket for one interface.

        Returns whether a responder socket was actually added.
        """
        # Join the group and create the responder via the same primitive
        # construction uses, so setup and rescan stay in lockstep. These are
        # user-initiated reconciles, so a requested interface that fails to
        # come up is surfaced once at warning (deduped per interface so the
        # polling monitor doesn't spam) rather than only at debug.
        respond_socket = add_interface(listen_socket, interface, apple_p2p=apple_p2p, unicast=self.zc.unicast)
        if respond_socket is None:
            self.zc.log_warning_once(f"Interface {interface!r} not added")
            return False
        try:
            await self._async_wrap_socket(respond_socket, is_sender=True)
        except Exception as exc:
            # Roll back the socket + group join on any failure so nothing is left
            # dangling.
            respond_socket.close()
            if listen_socket is not None:
                drop_multicast_member(listen_socket, interface)
            if not isinstance(exc, OSError):
                # Only an expected socket-level failure is best-effort; a real
                # bug (e.g. TypeError) must propagate rather than be downgraded
                # to a one-time "interface not added" warning.
                raise
            # Log and skip rather than abort the whole reconcile so the other
            # interfaces still come up (best-effort bring-up).
            self.zc.log_warning_once(f"Interface {interface!r} not added: {exc}")
            return False
        return True

    def _async_cache_cleanup(self) -> None:
        """Periodic cache cleanup."""
        now = current_time_millis()
        self.zc.question_history.async_expire(now)
        self.zc.record_manager.async_updates(
            now,
            [RecordUpdate(record, record) for record in self.zc.cache.async_expire(now)],
        )
        self.zc.record_manager.async_updates_complete(False)
        self._async_schedule_next_cache_cleanup()

    async def _async_close(self) -> None:
        """Cancel and wait for the cleanup task to finish."""
        assert self._setup_task is not None
        # Swallow CancelledError only if the setup task itself was
        # cancelled (close-before-start); outer-task cancellation must
        # propagate.
        try:
            await self._setup_task
        except asyncio.CancelledError:
            if not self._setup_task.cancelled():
                raise
        self._async_shutdown()
        await asyncio.sleep(0)  # flush out any call soons
        if self._cleanup_timer is not None:
            self._cleanup_timer.cancel()

    def _async_close_sender(self, wrapped: _WrappedTransport, listen_socket: socket.socket | None) -> None:
        """Drop a per-interface sender's wrappers/protocol and close its transport."""
        transport = wrapped.transport
        self._async_remove_transport(transport)
        try:
            if listen_socket is not None:
                drop_multicast_member(listen_socket, wrapped.multicast_interface)
        finally:
            # Release the socket even if a non-benign leave (e.g. EPERM) raises.
            transport.close()

    async def _async_create_endpoints(self) -> None:
        """Create endpoints to send and receive."""
        reader_sockets = []
        sender_sockets = []
        if self._listen_socket:
            reader_sockets.append(self._listen_socket)
        for s in self._respond_sockets:
            if s not in reader_sockets:
                reader_sockets.append(s)
            sender_sockets.append(s)

        for s in reader_sockets:
            reader = await self._async_wrap_socket(s, s in sender_sockets)
            # _async_wrap_socket registers the transport with no await between
            # creating and registering it, and the pending-handle cleanup below
            # adds no await either, so a concurrent shutdown always sees ``s``
            # in exactly one place.
            if s is self._listen_socket:
                # Keep a handle to the shared listen socket so interface
                # rescans can add/drop multicast memberships on it.
                self._listen_transport = reader
                self._listen_socket = None
            if s in self._respond_sockets:
                self._respond_sockets.remove(s)

    async def _async_rebuild_listen_socket(
        self,
        apple_p2p: bool,
        desired: dict[tuple[str, int], str | tuple[tuple[str, int, int], int]],
        current: dict[tuple[str, int], _WrappedTransport],
    ) -> None:
        """Replace the listen socket with one whose family covers the desired set.

        The listen socket's family is otherwise fixed at construction; this
        lets an instance start receiving a newly added address family, and is
        also used to convert a Default dual-use socket to a pure listener. The
        replacement family is derived from the desired set (not the
        requested ip_version, which an explicit list can contradict) so it
        always covers every desired interface and never needs an immediate
        re-rebuild. Interfaces that are staying are re-joined on the new socket,
        and the old socket is closed (releasing its memberships).
        """
        has_v6 = any(isinstance(interface, tuple) for interface in desired.values())
        has_v4 = any(not isinstance(interface, tuple) for interface in desired.values())
        if has_v4 and has_v6:
            family_version = IPVersion.All
        elif has_v6:
            family_version = IPVersion.V6Only
        else:
            family_version = IPVersion.V4Only
        new_listen = new_listen_socket(family_version, apple_p2p)
        if new_listen is None:
            raise RuntimeError("Failed to create a listen socket for the new interface family")
        try:
            for bind_address, interface in desired.items():
                # A staying interface that can't re-join on the new socket keeps
                # its sender but receives only via the shared socket it never
                # joined; surface that degraded state like _async_add_interface.
                if bind_address in current and not add_multicast_member(new_listen, interface):
                    self.zc.log_warning_once(
                        f"Interface {interface!r} could not re-join the multicast group "
                        "on the rebuilt listen socket"
                    )
            new_reader = await self._async_wrap_socket(new_listen, is_sender=False)
        except Exception:
            # Endpoint creation failed; close the unadopted socket (and its
            # joins) rather than leak it, mirroring _async_add_interface.
            new_listen.close()
            raise
        # A rebuild is only entered with a live listen socket, so the old
        # transport is always present.
        old_listen_transport = self._listen_transport
        assert old_listen_transport is not None
        self._listen_transport = new_reader
        old_transport = old_listen_transport.transport
        self._async_remove_transport(old_transport)
        old_transport.close()

    def _async_remove_transport(self, transport: asyncio.DatagramTransport) -> None:
        """Drop a transport's protocol/reader/sender wrappers, cancelling its timers."""
        kept_protocols = []
        for protocol in self.protocols:
            if protocol.transport is not None and protocol.transport.transport is transport:
                # Cancel any pending TC-reassembly timers so one can't fire a
                # response against the transport we're about to close.
                protocol.cancel_pending_timers()
            else:
                kept_protocols.append(protocol)
        self.protocols = kept_protocols
        self.readers = _without_transport(self.readers, transport)
        self.senders = _without_transport(self.senders, transport)

    def _async_schedule_next_cache_cleanup(self) -> None:
        """Schedule the next cache cleanup."""
        loop = self.loop
        assert loop is not None
        self._cleanup_timer = loop.call_at(loop.time() + _CACHE_CLEANUP_INTERVAL, self._async_cache_cleanup)

    async def _async_setup(self, loop_thread_ready: threading.Event | None) -> None:
        """Set up the instance."""
        self._async_schedule_next_cache_cleanup()
        await self._async_create_endpoints()
        assert self.running_future is not None
        if not self.running_future.done():
            self.running_future.set_result(True)
        if loop_thread_ready:
            loop_thread_ready.set()

    def _async_shutdown(self) -> None:
        """Shutdown transports and sockets; safe to call repeatedly."""
        assert self.running_future is not None
        assert self.loop is not None
        self.running_future = self.loop.create_future()
        # Cancel pending setup so it can't wrap fresh transports after
        # shutdown has started.
        if self._setup_task is not None and not self._setup_task.done():
            self._setup_task.cancel()
        for wrapped_transport in itertools.chain(self.senders, self.readers):
            wrapped_transport.transport.close()
        # Anything still here was never adopted by a transport.
        if self._listen_socket is not None:
            self._listen_socket.close()
            self._listen_socket = None
        for s in self._respond_sockets:
            s.close()
        self._respond_sockets = []

    async def _async_wrap_socket(self, sock: socket.socket, is_sender: bool) -> _WrappedTransport:
        """Adopt a socket into a transport, register it, and return the reader wrapper."""
        assert self.loop is not None
        transport, protocol = await self.loop.create_datagram_endpoint(  # type: ignore[type-var]
            lambda: AsyncListener(self.zc),  # type: ignore[arg-type, return-value]
            sock=sock,
        )
        datagram_transport = cast(asyncio.DatagramTransport, transport)
        reader = make_wrapped_transport(datagram_transport)
        # No ``await`` between wrapping and registering so a concurrent
        # shutdown always sees the transport in exactly one place.
        self.protocols.append(cast(AsyncListener, protocol))
        self.readers.append(reader)
        if is_sender:
            self.senders.append(make_wrapped_transport(datagram_transport))
        return reader