heos/heos.py

"""HEOS CLI client (TCP port 1255).

HEOS is not JSON-RPC or anything else standard: you send one
`heos://group/command?a=1&b=2` line and read back one JSON line. This
keeps the socket open between commands -- the panel fires a command on
every button press, and a fresh TCP handshake per press is what made
the original bridge feel sluggish.
"""

import json
import socket
import threading
import time
from urllib.parse import unquote_plus


class HeosError(RuntimeError):
    """The device answered, but said no (or never answered at all)."""


def parse_message(message: str) -> dict:
    """Turn a HEOS `heos.message` string into a dict.

    e.g. "pid=12345&level=23" -> {"pid": "12345", "level": "23"}
    """
    result = {}
    for part in message.split("&"):
        if "=" in part:
            key, value = part.split("=", 1)
            result[key] = unquote_plus(value)
        elif part:
            result[part] = ""
    return result


class HeosClient:
    """One persistent, lock-guarded, self-healing HEOS connection."""

    def __init__(self, host: str, port: int = 1255, timeout: float = 5.0):
        self.host = host
        self.port = port
        self.timeout = timeout
        self._sock = None
        self._buffer = b""
        self._lock = threading.Lock()
        # pid -> {"cur_pos", "duration"}, filled in from whatever
        # unsolicited progress events turn up while we are reading the
        # reply to some other command. Dropped on reconnect, since a gap
        # in the connection is a gap in what we know.
        self._progress = {}

    # -- connection management -----------------------------------------
    def _connect(self):
        self._close()
        sock = socket.create_connection((self.host, self.port), timeout=self.timeout)
        sock.settimeout(self.timeout)
        self._sock = sock
        self._buffer = b""
        self._progress = {}
        # Without this, HEOS never pushes the now-playing-progress events
        # that _exchange() below watches for. Best-effort: a player that
        # refuses it still works, it just never shows song progress.
        try:
            self._exchange("system/register_for_change_events", {"enable": "on"})
        except (OSError, HeosError):
            pass

    def _close(self):
        if self._sock is not None:
            try:
                self._sock.close()
            except OSError:
                pass
        self._sock = None
        self._buffer = b""

    def _read_line(self, deadline: float) -> bytes:
        while b"\r\n" not in self._buffer:
            remaining = deadline - time.monotonic()
            if remaining <= 0:
                raise TimeoutError("no reply from HEOS in time")
            self._sock.settimeout(remaining)
            chunk = self._sock.recv(4096)
            if not chunk:
                raise ConnectionError("HEOS closed the connection")
            self._buffer += chunk
        line, self._buffer = self._buffer.split(b"\r\n", 1)
        return line

    # -- the one method everything else goes through -------------------
    def command(self, path: str, **params) -> dict:
        """Send `heos://<path>?<params>` and return the parsed reply.

        Retries once on a socket-level problem, because HEOS quietly
        drops connections that have been idle for a few minutes.
        """
        with self._lock:
            for attempt in (1, 2):
                try:
                    if self._sock is None:
                        self._connect()
                    return self._exchange(path, params)
                except (OSError, ValueError) as exc:
                    # ValueError covers a desynced stream (bad JSON): in
                    # both cases the fix is the same, start a fresh socket.
                    self._close()
                    if attempt == 2:
                        raise HeosError(f"cannot reach HEOS at {self.host}: {exc}") from exc

    def _exchange(self, path: str, params: dict) -> dict:
        query = "&".join(f"{k}={v}" for k, v in params.items() if v is not None)
        command = f"heos://{path}" + (f"?{query}" if query else "")
        self._sock.sendall(command.encode("utf-8") + b"\r\n")

        # Browse calls can take a while, hence the generous overall budget.
        deadline = time.monotonic() + self.timeout * 3
        while True:
            reply = json.loads(self._read_line(deadline).decode("utf-8"))
            heos = reply.get("heos", {})

            if heos.get("command") != path:
                self._watch_progress(heos)
                continue  # an event or a late reply to something else
            if "under process" in heos.get("message", ""):
                continue  # placeholder ack; the real payload follows
            if heos.get("result") == "fail":
                raise HeosError(_failure_text(heos))
            return reply

    def command_batch(self, requests: list) -> list:
        """Send several commands back-to-back on the one socket, then
        collect their replies as they come in -- so N independent reads
        (a room's volume, play state and now-playing, say) cost one round
        trip's worth of latency instead of N.

        `requests` is a list of (path, params) pairs. Returns one entry
        per request, in the same order: the parsed reply dict, or the
        HeosError it failed with -- a single command failing does not
        sink the rest of the batch. Only a connection-level problem
        raises, same as command().
        """
        with self._lock:
            for attempt in (1, 2):
                try:
                    if self._sock is None:
                        self._connect()
                    return self._exchange_batch(requests)
                except (OSError, ValueError) as exc:
                    self._close()
                    if attempt == 2:
                        raise HeosError(f"cannot reach HEOS at {self.host}: {exc}") from exc

    def _exchange_batch(self, requests: list) -> list:
        for path, params in requests:
            query = "&".join(f"{k}={v}" for k, v in params.items() if v is not None)
            command = f"heos://{path}" + (f"?{query}" if query else "")
            self._sock.sendall(command.encode("utf-8") + b"\r\n")

        results = [None] * len(requests)
        pending = list(range(len(requests)))
        deadline = time.monotonic() + self.timeout * 3
        while pending:
            reply = json.loads(self._read_line(deadline).decode("utf-8"))
            heos = reply.get("heos", {})
            if "under process" in heos.get("message", ""):
                continue  # placeholder ack; the real payload follows

            index = self._claim(pending, requests, heos)
            if index is None:
                self._watch_progress(heos)
                continue  # an event, or a reply to a command outside this batch
            pending.remove(index)
            results[index] = HeosError(_failure_text(heos)) if heos.get("result") == "fail" else reply
        return results

    @staticmethod
    def _claim(pending, requests, heos):
        """Which pending request this reply answers. Matches on path
        first; when more than one pending request shares a path (e.g.
        get_now_playing_media for two different pids), whichever id
        parameter -- pid/gid/sid -- the request carried ties it to the
        right reply, since HEOS echoes it back in the message."""
        candidates = [i for i in pending if requests[i][0] == heos.get("command")]
        if not candidates:
            return None
        if len(candidates) == 1:
            return candidates[0]
        message = parse_message(heos.get("message", ""))
        for i in candidates:
            params = requests[i][1]
            for id_key in ("pid", "gid", "sid"):
                if id_key in params and message.get(id_key) == str(params[id_key]):
                    return i
        return candidates[0]  # can't tell them apart; oldest first

    def _watch_progress(self, heos: dict):
        """Skims a passing player_now_playing_progress event for its
        pid/cur_pos/duration, the only source of song-position data this
        client has -- there is no dedicated listener, so this only catches
        what happens to arrive while some other command is being read."""
        if heos.get("command") != "event/player_now_playing_progress":
            return
        message = parse_message(heos.get("message", ""))
        pid = message.get("pid")
        if not pid:
            return
        try:
            self._progress[pid] = {
                "cur_pos": int(message.get("cur_pos", 0)),
                "duration": int(message.get("duration", 0)),
            }
        except ValueError:
            pass

    def progress_for(self, pid) -> dict:
        """{"cur_pos", "duration"} in ms for a player, or None if no
        progress event for it has come through yet this connection."""
        return self._progress.get(str(pid))

    def heart_beat(self):
        """Keep the socket warm so the first press after an idle spell is
        as fast as the rest."""
        self.command("system/heart_beat")


def _failure_text(heos: dict) -> str:
    message = parse_message(heos.get("message", ""))
    text = message.get("text") or "unknown error"
    eid = message.get("eid")
    return f"HEOS refused '{heos.get('command')}': {text}" + (f" (eid {eid})" if eid else "")