"""Finding receivers on the network.

Three protocols, three mechanisms, one device list:
  * Google Cast — mDNS `_googlecast._tcp`, browsed by pychromecast.
  * AirPlay     — mDNS `_airplay._tcp`.
  * DLNA        — SSDP `MediaRenderer`, then the device description XML for its
                  control URLs.
"""
from __future__ import annotations

import asyncio
import contextlib
import ipaddress
import socket
import time
import xml.etree.ElementTree as ET
from urllib.parse import urljoin, urlsplit

import aiohttp
from zeroconf import ServiceBrowser, ServiceStateChange, Zeroconf

from .devices import AIRPLAY, CAST, DLNA, Device

# Caps on things a hostile or broken LAN can hand us. Every one of these is a
# number an attacker would otherwise get to choose.
MAX_SSDP_REPLIES = 32          # description URLs fetched per scan
DESCRIPTION_LIMIT = 256 << 10  # bytes of device XML we are willing to parse
MAX_DEVICES = 64               # entries in the device table


def _is_lan(host: str) -> bool:
    """True for an address a receiver could plausibly have."""
    try:
        ip = ipaddress.ip_address(host)
    except ValueError:
        return False
    return ip.is_private or ip.is_loopback or ip.is_link_local
from .util import log

AIRPLAY_TYPE = "_airplay._tcp.local."
GOOGLECAST_TYPE = "_googlecast._tcp.local."
SSDP_ADDR, SSDP_PORT = "239.255.255.250", 1900
SSDP_TARGET = "urn:schemas-upnp-org:device:MediaRenderer:1"
STALE_AFTER = 180.0     # a device unseen this long drops off the list
AWAY_GRACE = 60.0       # ...but a withdrawn announcement gets this long to come back


class Discovery:
    def __init__(self, on_change) -> None:
        self.devices: dict[str, Device] = {}
        self.on_change = on_change
        self.zeroconf: Zeroconf | None = None
        self.cast_browser = None
        self._airplay_browser: ServiceBrowser | None = None
        self._cast_txt_browser: ServiceBrowser | None = None
        self._loop: asyncio.AbstractEventLoop | None = None
        self._ssdp_task: asyncio.Task | None = None
        self._reap_task: asyncio.Task | None = None
        self._session: aiohttp.ClientSession | None = None

    # ---- lifecycle ----
    async def start(self) -> None:
        self._loop = asyncio.get_running_loop()
        self.zeroconf = Zeroconf()
        self._session = aiohttp.ClientSession(
            timeout=aiohttp.ClientTimeout(total=6), headers={"User-Agent": "screencast/1.0"}
        )
        self._start_cast()
        self._airplay_browser = ServiceBrowser(
            self.zeroconf, [AIRPLAY_TYPE], handlers=[self._airplay_change]
        )
        # Cast devices announce what they are showing in their TXT record, so the
        # list can say "Netflix" or "idle" without opening a connection to each.
        self._cast_txt_browser = ServiceBrowser(
            self.zeroconf, [GOOGLECAST_TYPE], handlers=[self._cast_txt_change]
        )
        self._ssdp_task = asyncio.create_task(self._ssdp_loop())
        self._reap_task = asyncio.create_task(self._reap_loop())

    def _start_cast(self) -> None:
        try:
            from pychromecast.discovery import CastBrowser, SimpleCastListener
        except Exception as exc:  # noqa: BLE001
            log.warning("Google Cast discovery unavailable: %s", exc)
            return
        listener = SimpleCastListener(
            add_callback=self._cast_seen, update_callback=self._cast_seen,
            remove_callback=self._cast_gone,
        )
        self.cast_browser = CastBrowser(listener, self.zeroconf)
        self.cast_browser.start_discovery()

    async def stop(self) -> None:
        for task in (self._ssdp_task, self._reap_task):
            if task:
                task.cancel()
        if self.cast_browser is not None:
            try:
                self.cast_browser.stop_discovery()
            except Exception:  # noqa: BLE001
                pass
        for browser in (self._airplay_browser, self._cast_txt_browser):
            if browser is not None:
                browser.cancel()
        if self.zeroconf is not None:
            # close() does blocking network I/O; keep it off the event loop.
            zc, self.zeroconf = self.zeroconf, None
            with contextlib.suppress(Exception):
                await asyncio.wait_for(asyncio.to_thread(zc.close), 3)
        if self._session is not None:
            await self._session.close()

    async def rescan(self) -> None:
        """Ask again now instead of waiting for the next announcement.

        The long-lived browsers are left alone (stopping pychromecast's browser
        tears down the shared Zeroconf instance with it); a throwaway browser
        over the same types is enough to put fresh queries on the wire.
        """
        if self.zeroconf is not None:
            probe = ServiceBrowser(self.zeroconf, [GOOGLECAST_TYPE, AIRPLAY_TYPE],
                                   handlers=[self._probe_change])
            asyncio.get_running_loop().call_later(4, probe.cancel)
        await self._ssdp_search()

    def _probe_change(self, zeroconf: Zeroconf, service_type: str, name: str,
                      state_change: ServiceStateChange) -> None:
        """A manual scan only needs to make devices answer; the real browsers record them."""
        return

    # ---- Google Cast (callbacks arrive on a zeroconf thread) ----
    def _cast_seen(self, uuid, service) -> None:
        if self._loop:
            self._loop.call_soon_threadsafe(self._cast_seen_main, uuid)

    def _cast_gone(self, uuid, service, cast_info) -> None:
        if self._loop:
            self._loop.call_soon_threadsafe(self._drop, f"cast:{uuid}")

    def _cast_seen_main(self, uuid) -> None:
        info = (self.cast_browser.devices or {}).get(uuid) if self.cast_browser else None
        if info is None:
            return
        dev = self._upsert(
            Device(
                id=f"cast:{uuid}",
                kind=CAST,
                name=info.friendly_name or str(uuid),
                host=info.host,
                port=info.port or 8009,
                model=info.model_name or "",
                manufacturer=info.manufacturer or "",
                extra={"uuid": str(uuid), "cast_type": info.cast_type},
            )
        )
        dev.extra["cast_info"] = info

    def _cast_txt_change(self, zeroconf: Zeroconf, service_type: str, name: str,
                         state_change: ServiceStateChange) -> None:
        if self._loop is None or state_change is ServiceStateChange.Removed:
            return
        info = zeroconf.get_service_info(service_type, name, timeout=2000)
        if info is None:
            return
        txt = {
            k.decode(errors="replace"): (v or b"").decode(errors="replace")
            for k, v in (info.properties or {}).items()
        }
        raw = txt.get("id", "")
        if len(raw) != 32:
            return
        uuid = f"{raw[0:8]}-{raw[8:12]}-{raw[12:16]}-{raw[16:20]}-{raw[20:]}"
        self._loop.call_soon_threadsafe(self._cast_txt_main, uuid, txt)

    def _cast_txt_main(self, uuid: str, txt: dict) -> None:
        dev = self.devices.get(f"cast:{uuid}")
        if dev is None:
            return
        app = txt.get("rs", "").strip()
        # st: 0 = nothing running, anything else = an app has the screen.
        status = "idle" if txt.get("st", "0") == "0" else "busy"
        if (dev.app, dev.status) != (app, status):
            dev.app, dev.status = app, status
            self.on_change()

    # ---- AirPlay ----
    def _airplay_change(self, zeroconf: Zeroconf, service_type: str, name: str,
                        state_change: ServiceStateChange) -> None:
        if self._loop is None:
            return
        if state_change is ServiceStateChange.Removed:
            # TVs withdraw and re-publish their AirPlay record constantly, and a
            # row that disappears under the pointer is worse than a stale one:
            # mark it away and let the reaper drop it if it stays gone.
            self._loop.call_soon_threadsafe(self._mark_away, f"airplay:{name}")
            return
        info = zeroconf.get_service_info(service_type, name, timeout=2000)
        if info is None:
            return
        addrs = info.parsed_scoped_addresses() if hasattr(info, "parsed_scoped_addresses") else []
        ipv4 = next((a for a in addrs if ":" not in a), None) or (addrs[0] if addrs else "")
        if not ipv4:
            return
        txt = {
            k.decode(errors="replace"): (v or b"").decode(errors="replace")
            for k, v in (info.properties or {}).items()
        }
        friendly = name.split(".")[0].replace("\\032", " ")
        self._loop.call_soon_threadsafe(
            self._upsert,
            Device(
                id=f"airplay:{name}",
                kind=AIRPLAY,
                name=txt.get("name") or friendly,
                host=ipv4,
                port=info.port or 7000,
                model=txt.get("model", ""),
                manufacturer="Apple" if txt.get("model", "").startswith(("Mac", "Apple")) else "",
                extra={"txt": txt, "service": name},
            ),
        )

    # ---- DLNA ----
    async def _ssdp_loop(self) -> None:
        while True:
            try:
                await self._ssdp_search()
            except asyncio.CancelledError:
                raise
            except Exception as exc:  # noqa: BLE001
                log.debug("ssdp: %s", exc)
            await asyncio.sleep(60)

    async def _ssdp_search(self) -> None:
        """M-SEARCH from every local address; renderers answer by unicast."""
        msg = (
            "M-SEARCH * HTTP/1.1\r\n"
            f"HOST: {SSDP_ADDR}:{SSDP_PORT}\r\n"
            'MAN: "ssdp:discover"\r\n'
            "MX: 2\r\n"
            f"ST: {SSDP_TARGET}\r\n\r\n"
        ).encode()
        locations: dict[str, str] = {}   # description URL -> the address that sent it
        for local_ip in _local_ipv4s():
            sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            sock.setsockopt(socket.IPPROTO_IP, socket.IP_MULTICAST_TTL, 2)
            try:
                sock.bind((local_ip, 0))
                sock.setblocking(False)
                loop = asyncio.get_running_loop()
                await loop.sock_sendto(sock, msg, (SSDP_ADDR, SSDP_PORT))
                deadline = time.monotonic() + 3
                while time.monotonic() < deadline:
                    try:
                        data, src = await asyncio.wait_for(
                            loop.sock_recvfrom(sock, 4096), max(0.2, deadline - time.monotonic())
                        )
                    except (asyncio.TimeoutError, OSError):
                        break
                    loc = _header(data.decode(errors="replace"), "location")
                    if loc and len(locations) < MAX_SSDP_REPLIES:
                        locations.setdefault(loc, src[0])
            except OSError as exc:
                log.debug("ssdp on %s: %s", local_ip, exc)
            finally:
                sock.close()
        for loc, sender in locations.items():
            await self._add_dlna(loc, sender)

    async def _add_dlna(self, location: str, sender: str) -> None:
        assert self._session is not None
        # An SSDP reply is an unauthenticated UDP packet from anyone on the wire,
        # and we are about to fetch the URL inside it as this user. Only fetch it
        # back from the host that sent it: otherwise the packet chooses the
        # target, and "http://127.0.0.1:9091/..." is a request we would make for
        # a stranger.
        host = urlsplit(location).hostname or ""
        if host != sender or not _is_lan(host):
            log.info("ssdp: ignoring %s announced by %s", location, sender)
            return
        try:
            async with self._session.get(location) as resp:
                # No renderer needs a megabyte to describe itself; an attacker
                # would happily send us gigabytes, or an entity bomb.
                raw = await resp.content.read(DESCRIPTION_LIMIT)
        except Exception as exc:  # noqa: BLE001
            log.debug("dlna description %s: %s", location, exc)
            return
        xml = raw.decode("utf-8", errors="replace")
        try:
            root = ET.fromstring(xml)
        except ET.ParseError:
            return
        ns = {"u": "urn:schemas-upnp-org:device-1-0"}
        dev_el = root.find("u:device", ns)
        if dev_el is None:
            return
        udn = (dev_el.findtext("u:UDN", "", ns) or location).strip()
        name = (dev_el.findtext("u:friendlyName", "", ns) or "DLNA renderer").strip()
        model = (dev_el.findtext("u:modelName", "", ns) or "").strip()
        maker = (dev_el.findtext("u:manufacturer", "", ns) or "").strip()
        base = root.findtext("u:URLBase", "", ns) or location
        control: dict[str, str] = {}
        for svc in dev_el.iter("{urn:schemas-upnp-org:device-1-0}service"):
            stype = svc.findtext("u:serviceType", "", ns)
            curl = svc.findtext("u:controlURL", "", ns)
            if not stype or not curl:
                continue
            target = urljoin(base, curl)
            # URLBase and controlURL are the device's own words. An absolute URL
            # here would point our POSTs anywhere it likes.
            if urlsplit(target).hostname != host:
                log.info("dlna: %s points its control URL at %s, ignoring", location, target)
                continue
            if "AVTransport" in stype:
                control["avtransport"] = target
                control["avtransport_type"] = stype
            elif "RenderingControl" in stype:
                control["rendering"] = target
                control["rendering_type"] = stype
        if "avtransport" not in control:
            return   # can be discovered but cannot be told to play anything
        self._upsert(
            Device(
                id=f"dlna:{udn}",
                kind=DLNA,
                name=name,
                host=host,
                port=int(location.split("/")[2].split(":")[1]) if ":" in location.split("/")[2] else 80,
                model=model,
                manufacturer=maker,
                extra={"control": control, "location": location},
            )
        )

    # ---- list bookkeeping ----
    def _upsert(self, dev: Device) -> Device:
        old = self.devices.get(dev.id)
        if old is not None:
            old.seen = time.time()
            changed = (old.name, old.host, old.port) != (dev.name, dev.host, dev.port)
            if old.status == "away":
                old.status, changed = "", True
            old.name, old.host, old.port = dev.name, dev.host, dev.port
            old.model = dev.model or old.model
            old.extra.update(dev.extra)
            if changed:
                self.on_change()
            return old
        if len(self.devices) >= MAX_DEVICES:
            # Announcements are free to send and we push the whole list to the
            # shell on every change; a flood must not become the shell's problem.
            log.warning("device list full (%d), ignoring %s", MAX_DEVICES, dev.id)
            return dev
        self.devices[dev.id] = dev
        log.info("found %s: %s (%s)", dev.kind, dev.name, dev.host)
        self.on_change()
        return dev

    def _mark_away(self, device_id: str) -> None:
        dev = self.devices.get(device_id)
        if dev is None or dev.status == "away":
            return
        dev.status = "away"
        dev.seen = time.time() - (STALE_AFTER - AWAY_GRACE)
        log.info("away %s", device_id)
        self.on_change()

    def _drop(self, device_id: str) -> None:
        if self.devices.pop(device_id, None) is not None:
            log.info("lost %s", device_id)
            self.on_change()

    async def _reap_loop(self) -> None:
        while True:
            await asyncio.sleep(15)
            now = time.time()
            for dev_id, dev in list(self.devices.items()):
                # Cast devices are kept by their browser; the rest age out —
                # AirPlay only once its away grace has run down too.
                if dev.kind in (DLNA, AIRPLAY) and now - dev.seen > STALE_AFTER:
                    self._drop(dev_id)


def _header(text: str, name: str) -> str:
    for line in text.splitlines():
        if line.lower().startswith(name + ":"):
            return line.split(":", 1)[1].strip()
    return ""


def _local_ipv4s() -> list[str]:
    """Every non-loopback IPv4 we can send from (several NICs, VPNs, …)."""
    out: list[str] = []
    try:
        import ifaddr

        for adapter in ifaddr.get_adapters():
            for ip in adapter.ips:
                if ip.is_IPv4 and not str(ip.ip).startswith("127."):
                    out.append(str(ip.ip))
    except Exception:  # noqa: BLE001
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        try:
            s.connect(("8.8.8.8", 9))
            out.append(s.getsockname()[0])
        except OSError:
            pass
        finally:
            s.close()
    return out or ["0.0.0.0"]