#!/usr/bin/python3
#
# Copyright (C) 2026 Danilo M. <danix@danix.xyz>
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 2 as
# published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# The headphones panel's backend: noise cancelling on Sony headphones that
# speak the "v2" MDR protocol (the WH-CH720N here) over RFCOMM. Volume and
# battery are not here: volume is AVRCP through PipeWire, and battery is
# BlueZ's Battery1, both of which the island already reads.
#
#   sonyctl MAC status               {ok, mode, level, voice}
#   sonyctl MAC nc|off               set the mode, print the new state
#   sonyctl MAC ambient [LEVEL]      ambient sound, LEVEL 1-20 (default: kept)
#   sonyctl MAC voice on|off         ambient sound with focus on voice on or off
#
# Every set keeps the level and focus on voice it does not change.
#   sonyctl MAC poweroff             turn the headphones off
#
# mode is "nc", "ambient" or "off". Errors print {ok: false, error}.
#
# Frames, as measured on a WH-CH720N: 3e, type, seq, 4-byte big-endian
# length, payload, 1-byte sum of everything after 3e, 3c. 3c/3d/3e inside a
# frame are sent as 3d then the byte & 0xef. Type 0c is a command, 01 an
# ack; each side acks with seq 1 - the received seq, and the next command
# uses the seq of the last ack received. The headphones answer nothing until
# an init (00 00) has been sent on the connection.
#
# The ambient payload subtype is 17 on this model. 15, which Gadgetbridge
# sends for headphones without wind-noise reduction, is acked and never
# answered. The state is 17, 01, on, ambient, focus-on-voice, level.

import errno
import json
import socket
import sys
import time

CMD, ACK = 0x0C, 0x01
# ponytail: the SDP record measured on a WH-CH720N puts the MDR service
# (956c7b26-...) on channel 18. Look it up over SDP if another model differs.
CHANNEL = 18
SPECIAL = (0x3C, 0x3D, 0x3E)


def encode(kind, seq, payload):
    body = bytes([kind, seq]) + len(payload).to_bytes(4, "big") + bytes(payload)
    body += bytes([sum(body) & 0xFF])
    out = bytearray([0x3E])
    for b in body:
        out += bytes([0x3D, b & 0xEF]) if b in SPECIAL else bytes([b])
    return bytes(out) + b"\x3c"


def decode(raw):
    """One frame, without its 3c trailer, to (type, seq, payload) or None."""
    if not raw.startswith(b"\x3e"):
        return None
    body, i = bytearray(), 1
    while i < len(raw):
        if raw[i] == 0x3D and i + 1 < len(raw):
            i += 1
            body.append(raw[i] | 0x10)
        else:
            body.append(raw[i])
        i += 1
    if len(body) < 7 or sum(body[:-1]) & 0xFF != body[-1]:
        return None
    n = int.from_bytes(body[2:6], "big")
    if n != len(body) - 7:
        return None
    return body[0], body[1], bytes(body[6:-1])


def state(p):
    """A 67/69 17 payload to {mode, level, voice}."""
    mode = "off" if p[3] == 0 else ("ambient" if p[4] else "nc")
    return {"mode": mode, "level": p[6], "voice": p[5] == 1}


class Link:
    def __init__(self, mac):
        # A connect right after the previous run closed fails with EBUSY
        # (measured) while the old channel is still going down.
        for attempt in range(10):
            self.s = socket.socket(socket.AF_BLUETOOTH, socket.SOCK_STREAM, socket.BTPROTO_RFCOMM)
            self.s.settimeout(3)
            try:
                self.s.connect((mac, CHANNEL))
                break
            except OSError as e:
                self.s.close()
                if e.errno != errno.EBUSY or attempt == 9:
                    raise
                time.sleep(0.3)
        self.seq = 0
        self.buf = b""

    def frames(self):
        while True:
            while b"\x3c" in self.buf:
                raw, self.buf = self.buf.split(b"\x3c", 1)
                f = decode(raw[raw.find(b"\x3e"):]) if b"\x3e" in raw else None
                if f:
                    yield f
            chunk = self.s.recv(1024)
            if not chunk:
                raise ConnectionError("headphones closed the connection")
            self.buf += chunk

    def request(self, payload, reply):
        """Send a command and return the first command payload starting with
        one of the reply codes, acking everything that arrives meanwhile."""
        self.s.send(encode(CMD, self.seq, payload))
        for kind, seq, p in self.frames():
            if kind == ACK:
                self.seq = seq
                continue
            self.s.send(encode(ACK, 1 - seq, b""))
            if p and p[0] in reply:
                return p


def run(mac, cmd, arg):
    link = Link(mac)
    try:
        link.request([0x00, 0x00], (0x01,))
        if cmd == "poweroff":
            # Unanswered: the headphones go away instead. The ack is enough.
            link.s.send(encode(CMD, link.seq, [0x24, 0x03, 0x01]))
            for kind, _, _ in link.frames():
                if kind == ACK:
                    return {"ok": True}
        cur = state(link.request([0x66, 0x17], (0x67,)))
        if cmd == "status":
            return {"ok": True, **cur}
        level, voice = cur["level"], int(cur["voice"])
        if cmd == "voice":
            voice = arg
        elif arg is not None:
            level = arg
        on = 0 if cmd == "off" else 1
        amb = 1 if cmd in ("ambient", "voice") else 0
        p = link.request([0x68, 0x17, 0x01, on, amb, voice, level], (0x69,))
        return {"ok": True, **state(p)}
    finally:
        link.s.close()


def main(argv):
    if len(argv) < 3 or argv[2] not in ("status", "nc", "ambient", "voice", "off", "poweroff"):
        return {"ok": False, "error": "usage: sonyctl MAC status|nc|ambient [LEVEL]|voice on|off|off|poweroff"}
    arg = None
    if argv[2] == "voice":
        if len(argv) < 4 or argv[3] not in ("on", "off"):
            return {"ok": False, "error": "voice takes on or off"}
        arg = int(argv[3] == "on")
    if argv[2] == "ambient" and len(argv) > 3:
        if not argv[3].isdigit() or not 1 <= int(argv[3]) <= 20:
            return {"ok": False, "error": "level must be 1-20"}
        arg = int(argv[3])
    try:
        return run(argv[1], argv[2], arg)
    except (OSError, ConnectionError) as e:
        return {"ok": False, "error": str(e) or type(e).__name__}


if __name__ == "__main__":
    r = main(sys.argv)
    print(json.dumps(r))
    sys.exit(0 if r["ok"] else 1)
