"""Native macOS system-audio capture via Core Audio process taps.

macOS 14.2 added *process taps* to the HAL: a process may ask ``coreaudiod``
for a copy of the audio that is being rendered, without a kernel extension, a
virtual loopback (BlackHole, Loopback, Soundflower, …), a Multi-Output Device
or a helper binary.  The tap is driven here straight through Core Audio's C
API with :mod:`ctypes`, so the only dependency is macOS itself.

Capture is a three step dance, all of it native:

1. a ``CATapDescription`` (global stereo mix-down, unmuted) is turned into a
   tap object with ``AudioHardwareCreateProcessTap``;
2. the tap is wrapped in a *private* aggregate device so it can be read like
   any other input device;
3. an IOProc on that aggregate device delivers interleaved ``float32`` frames
   from the HAL's real-time thread, which are buffered for the analyser.

Tapping other applications' audio is privacy sensitive: the first attempt
makes macOS ask for **System Audio Recording** permission (System Settings ▸
Privacy & Security).  Until that is granted the tap is created successfully
but delivers digital silence, so callers must treat a silent tap as a failed
capture.

The module is importable everywhere; :func:`available` is ``False`` off macOS
or on systems older than the tap API.
"""

from __future__ import annotations

import ctypes
import ctypes.util
import sys
import threading
import time
import uuid
from collections import deque
from typing import Any

__all__ = ["SystemAudioTap", "TapError", "available"]

#: ``kAudioObjectPropertyScopeGlobal`` / ``kAudioObjectPropertyElementMain``.
_SCOPE_GLOBAL = int.from_bytes(b"glob", "big")
_ELEMENT_MAIN = 0

#: Tap object property selectors (``CoreAudio/AudioHardware.h``).
_TAP_UID = int.from_bytes(b"tuid", "big")
_TAP_FORMAT = int.from_bytes(b"tfmt", "big")

#: ``kAudioFormatFlagIsFloat | kAudioFormatFlagIsPacked``.
_FORMAT_LPCM = int.from_bytes(b"lpcm", "big")
_FLAG_FLOAT = 0x1

#: Aggregate-device description keys (``CoreAudio/AudioHardware.h``).
_AGGREGATE_UID = "uid"
_AGGREGATE_NAME = "name"
_AGGREGATE_PRIVATE = "private"
_AGGREGATE_STACKED = "stacked"
_AGGREGATE_TAPS = "taps"
_AGGREGATE_TAP_AUTOSTART = "tapautostart"
_SUB_TAP_UID = "uid"
_SUB_TAP_DRIFT = "drift"

_COREAUDIO = "/System/Library/Frameworks/CoreAudio.framework/CoreAudio"
_COREFOUNDATION = "/System/Library/Frameworks/CoreFoundation.framework/CoreFoundation"
_UTF8 = 0x08000100

#: Seconds of audio kept in the hand-off buffer before the oldest is dropped.
_BUFFER_SECONDS = 1.0


class TapError(RuntimeError):
    """The native audio tap could not be created."""


class _PropertyAddress(ctypes.Structure):
    _fields_ = [
        ("mSelector", ctypes.c_uint32),
        ("mScope", ctypes.c_uint32),
        ("mElement", ctypes.c_uint32),
    ]


class _StreamFormat(ctypes.Structure):
    """``AudioStreamBasicDescription``."""

    _fields_ = [
        ("mSampleRate", ctypes.c_double),
        ("mFormatID", ctypes.c_uint32),
        ("mFormatFlags", ctypes.c_uint32),
        ("mBytesPerPacket", ctypes.c_uint32),
        ("mFramesPerPacket", ctypes.c_uint32),
        ("mBytesPerFrame", ctypes.c_uint32),
        ("mChannelsPerFrame", ctypes.c_uint32),
        ("mBitsPerChannel", ctypes.c_uint32),
        ("mReserved", ctypes.c_uint32),
    ]


class _AudioBuffer(ctypes.Structure):
    _fields_ = [
        ("mNumberChannels", ctypes.c_uint32),
        ("mDataByteSize", ctypes.c_uint32),
        ("mData", ctypes.c_void_p),
    ]


class _AudioBufferList(ctypes.Structure):
    # The HAL hands us a variably sized struct; mNumberBuffers says how many
    # entries are valid.  The array is oversized so indexing stays in bounds
    # for the (few) buffers a tap aggregate ever reports.
    _fields_ = [
        ("mNumberBuffers", ctypes.c_uint32),
        ("mBuffers", _AudioBuffer * 8),
    ]


_IOProc = ctypes.CFUNCTYPE(
    ctypes.c_int32,
    ctypes.c_uint32,  # inDevice
    ctypes.c_void_p,  # inNow
    ctypes.POINTER(_AudioBufferList),  # inInputData
    ctypes.c_void_p,  # inInputTime
    ctypes.c_void_p,  # outOutputData
    ctypes.c_void_p,  # inOutputTime
    ctypes.c_void_p,  # inClientData
)

_loaded: dict[str, Any] = {}
_load_lock = threading.Lock()


def _frameworks() -> tuple[ctypes.CDLL, ctypes.CDLL, ctypes.CDLL]:
    """Load (and cache) CoreAudio, CoreFoundation and the Objective-C runtime."""
    with _load_lock:
        if not _loaded:
            ca = ctypes.CDLL(_COREAUDIO)
            cf = ctypes.CDLL(_COREFOUNDATION)
            objc = ctypes.CDLL(ctypes.util.find_library("objc") or "libobjc.A.dylib")
            _configure(ca, cf, objc)
            _loaded.update(ca=ca, cf=cf, objc=objc)
    return _loaded["ca"], _loaded["cf"], _loaded["objc"]


def _signature(lib: ctypes.CDLL, name: str, arguments: list, result: Any) -> None:
    """Declare a C function, skipping it when this macOS does not have it."""
    function = getattr(lib, name, None)
    if function is None:
        return
    function.argtypes = arguments
    function.restype = result


def _configure(ca: ctypes.CDLL, cf: ctypes.CDLL, objc: ctypes.CDLL) -> None:
    _signature(
        ca,
        "AudioObjectGetPropertyData",
        [
            ctypes.c_uint32,
            ctypes.POINTER(_PropertyAddress),
            ctypes.c_uint32,
            ctypes.c_void_p,
            ctypes.POINTER(ctypes.c_uint32),
            ctypes.c_void_p,
        ],
        ctypes.c_int32,
    )
    _signature(
        ca,
        "AudioHardwareCreateProcessTap",
        [ctypes.c_void_p, ctypes.POINTER(ctypes.c_uint32)],
        ctypes.c_int32,
    )
    _signature(ca, "AudioHardwareDestroyProcessTap", [ctypes.c_uint32], ctypes.c_int32)
    _signature(
        ca,
        "AudioHardwareCreateAggregateDevice",
        [ctypes.c_void_p, ctypes.POINTER(ctypes.c_uint32)],
        ctypes.c_int32,
    )
    _signature(
        ca, "AudioHardwareDestroyAggregateDevice", [ctypes.c_uint32], ctypes.c_int32
    )
    _signature(
        ca,
        "AudioDeviceCreateIOProcID",
        [ctypes.c_uint32, _IOProc, ctypes.c_void_p, ctypes.POINTER(ctypes.c_void_p)],
        ctypes.c_int32,
    )
    _signature(ca, "AudioDeviceDestroyIOProcID", [ctypes.c_uint32, ctypes.c_void_p], ctypes.c_int32)
    _signature(ca, "AudioDeviceStart", [ctypes.c_uint32, ctypes.c_void_p], ctypes.c_int32)
    _signature(ca, "AudioDeviceStop", [ctypes.c_uint32, ctypes.c_void_p], ctypes.c_int32)

    _signature(
        cf,
        "CFStringCreateWithCString",
        [ctypes.c_void_p, ctypes.c_char_p, ctypes.c_uint32],
        ctypes.c_void_p,
    )
    _signature(
        cf,
        "CFDictionaryCreateMutable",
        [ctypes.c_void_p, ctypes.c_long, ctypes.c_void_p, ctypes.c_void_p],
        ctypes.c_void_p,
    )
    _signature(
        cf,
        "CFDictionarySetValue",
        [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p],
        None,
    )
    _signature(
        cf,
        "CFArrayCreate",
        [ctypes.c_void_p, ctypes.POINTER(ctypes.c_void_p), ctypes.c_long, ctypes.c_void_p],
        ctypes.c_void_p,
    )
    _signature(cf, "CFRelease", [ctypes.c_void_p], None)

    _signature(objc, "objc_getClass", [ctypes.c_char_p], ctypes.c_void_p)
    _signature(objc, "sel_registerName", [ctypes.c_char_p], ctypes.c_void_p)


def _sender(result: Any, arguments: tuple[Any, ...]):
    """A typed ``objc_msgSend`` for one method signature (arm64 needs this)."""
    key = f"msg:{result}:{arguments}"
    objc = _frameworks()[2]
    with _load_lock:
        sender = _loaded.get(key)
        if sender is None:
            prototype = ctypes.CFUNCTYPE(
                result, ctypes.c_void_p, ctypes.c_void_p, *arguments
            )
            sender = prototype(("objc_msgSend", objc))
            _loaded[key] = sender
    return sender


def _selector(name: bytes) -> int:
    return int(_frameworks()[2].sel_registerName(name) or 0)


def _send_object(receiver: int, name: bytes, *arguments: int) -> int:
    sender = _sender(ctypes.c_void_p, (ctypes.c_void_p,) * len(arguments))
    return int(sender(receiver, _selector(name), *arguments) or 0)


def _send_flag(receiver: int, name: bytes, value: bool) -> None:
    _sender(None, (ctypes.c_bool,))(receiver, _selector(name), value)


def _send_integer(receiver: int, name: bytes, value: int) -> None:
    _sender(None, (ctypes.c_long,))(receiver, _selector(name), value)


def _osstatus(code: int) -> str:
    """Render an ``OSStatus`` as its four-character code when it has a printable one."""
    raw = (int(code) & 0xFFFFFFFF).to_bytes(4, "big")
    if all(32 <= byte < 127 for byte in raw):
        return f"{raw.decode()!r}"
    return str(int(code))


def available() -> bool:
    """True when this machine can create Core Audio process taps."""
    if sys.platform != "darwin":
        return False
    try:
        ca = _frameworks()[0]
    except OSError:
        return False
    return hasattr(ca, "AudioHardwareCreateProcessTap")


class SystemAudioTap:
    """Records what the default output device is playing (no driver required).

    ``start()`` raises :class:`TapError` when the tap cannot be created.
    Buffers arrive from the HAL as ``float32``; :meth:`read` hands them to the
    caller as raw little-endian PCM bytes.
    """

    def __init__(self) -> None:
        self.sample_rate = 48000
        self.channels = 2
        self.bytes_per_frame = 8
        self._tag = str(uuid.uuid4()).upper()
        self._tap_id: int | None = None
        self._device_id: int | None = None
        self._io_proc: ctypes.c_void_p | None = None
        self._callback: Any = None
        self._buffers: deque[bytes] = deque()
        self._buffered = 0
        self._limit = int(self.sample_rate * self.bytes_per_frame * _BUFFER_SECONDS)
        self._running = False
        self._cond = threading.Condition()

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

    def start(self) -> None:
        if self._running:
            return
        if not available():
            raise TapError("Core Audio process taps are unavailable on this system")

        ca, cf, _ = _frameworks()
        owned: list[int] = []

        def cfstr(text: str) -> int:
            ref = int(cf.CFStringCreateWithCString(None, text.encode(), _UTF8) or 0)
            if not ref:
                raise TapError("could not allocate a Core Foundation string")
            owned.append(ref)
            return ref

        try:
            description, nsuuid = self._describe(cfstr)
            owned.extend((description, nsuuid))

            tap_id = ctypes.c_uint32(0)
            status = ca.AudioHardwareCreateProcessTap(description, ctypes.byref(tap_id))
            if status != 0:
                raise TapError(
                    f"AudioHardwareCreateProcessTap failed ({_osstatus(status)})"
                )
            self._tap_id = tap_id.value

            self._read_format(ca)
            uid = self._read_uid(ca)
            if not uid:
                raise TapError("could not read the tap UID")
            owned.append(uid)

            self._build_device(ca, cf, cfstr, uid, owned)
            self._start_io(ca)
        except Exception:
            self._teardown()
            raise
        finally:
            for ref in owned:
                cf.CFRelease(ref)

        with self._cond:
            self._running = True
            self._cond.notify_all()

    def stop(self) -> None:
        self._teardown()
        with self._cond:
            self._running = False
            self._buffers.clear()
            self._buffered = 0
            self._cond.notify_all()

    @property
    def running(self) -> bool:
        return self._running

    # ----- capture ------------------------------------------------------ #

    def read(self, timeout: float | None = 1.0) -> bytes | None:
        """Return the next buffer of PCM, or ``None`` on timeout / after stop."""
        deadline = None if timeout is None else time.monotonic() + timeout
        with self._cond:
            while not self._buffers:
                if not self._running:
                    return None
                if deadline is None:
                    self._cond.wait(0.5)
                    continue
                remaining = deadline - time.monotonic()
                if remaining <= 0:
                    return None
                self._cond.wait(remaining)
            data = self._buffers.popleft()
            self._buffered -= len(data)
            return data

    # ----- Core Audio plumbing ------------------------------------------ #

    def _describe(self, cfstr) -> tuple[int, int]:
        """Build the ``CATapDescription`` for a global, unmuted stereo tap."""
        objc = _frameworks()[2]
        cls = objc.objc_getClass(b"CATapDescription")
        if not cls:
            raise TapError("CATapDescription is unavailable on this system")

        # "…but exclude these processes": an empty list means everything the
        # output device plays.
        empty = _empty_array()
        try:
            description = _send_object(
                _send_object(int(cls), b"alloc"),
                b"initStereoGlobalTapButExcludeProcesses:",
                empty,
            )
        finally:
            # The new CATapDescription has copied the array.
            _frameworks()[1].CFRelease(empty)
        if not description:
            raise TapError("could not create a CATapDescription")

        nsuuid = _send_object(
            _send_object(int(objc.objc_getClass(b"NSUUID")), b"alloc"),
            b"initWithUUIDString:",
            cfstr(self._tag),
        )
        if not nsuuid:
            raise TapError("could not create the tap UUID")
        _send_object(description, b"setUUID:", nsuuid)
        _send_object(description, b"setName:", cfstr("beefy-player"))
        _send_flag(description, b"setPrivate:", True)
        # 0 == CATapUnmuted: the audio we record keeps playing as usual.
        _send_integer(description, b"setMuteBehavior:", 0)
        return description, nsuuid

    def _read_uid(self, ca: ctypes.CDLL) -> int:
        uid = ctypes.c_void_p(0)
        size = ctypes.c_uint32(ctypes.sizeof(ctypes.c_void_p))
        if _get_property(ca, self._tap_id, _TAP_UID, uid, size) != 0:
            return 0
        return int(uid.value or 0)

    def _read_format(self, ca: ctypes.CDLL) -> None:
        fmt = _StreamFormat()
        size = ctypes.c_uint32(ctypes.sizeof(fmt))
        if _get_property(ca, self._tap_id, _TAP_FORMAT, fmt, size) != 0:
            raise TapError("could not read the tap's stream format")
        if (
            fmt.mFormatID != _FORMAT_LPCM
            or fmt.mBitsPerChannel != 32
            or not fmt.mFormatFlags & _FLAG_FLOAT
        ):
            raise TapError(
                "unsupported tap format "
                f"({fmt.mFormatID.to_bytes(4, 'big')!r}, {fmt.mBitsPerChannel} bit)"
            )
        self.sample_rate = round(fmt.mSampleRate) or 48000
        self.channels = int(fmt.mChannelsPerFrame) or 1
        self.bytes_per_frame = int(fmt.mBytesPerFrame or 4 * self.channels)
        self._limit = int(self.sample_rate * self.bytes_per_frame * _BUFFER_SECONDS)

    def _build_device(
        self, ca: ctypes.CDLL, cf: ctypes.CDLL, cfstr, uid: int, owned: list[int]
    ) -> None:
        tap_entry = int(cf.CFDictionaryCreateMutable(None, 0, None, None) or 0)
        if not tap_entry:
            raise TapError("could not allocate the tap description dictionary")
        owned.append(tap_entry)
        cf.CFDictionarySetValue(tap_entry, cfstr(_SUB_TAP_UID), uid)
        cf.CFDictionarySetValue(tap_entry, cfstr(_SUB_TAP_DRIFT), _boolean(True))
        entries = (ctypes.c_void_p * 1)(tap_entry)
        tap_list = int(cf.CFArrayCreate(None, entries, 1, None) or 0)
        if not tap_list:
            raise TapError("could not allocate the tap list")
        owned.append(tap_list)

        aggregate = int(cf.CFDictionaryCreateMutable(None, 0, None, None) or 0)
        if not aggregate:
            raise TapError("could not allocate the aggregate description")
        owned.append(aggregate)
        for key, value in (
            (_AGGREGATE_UID, cfstr(str(uuid.uuid4()).upper())),
            (_AGGREGATE_NAME, cfstr("beefy-player tap")),
            (_AGGREGATE_PRIVATE, _boolean(True)),
            (_AGGREGATE_STACKED, _boolean(False)),
            (_AGGREGATE_TAPS, tap_list),
            (_AGGREGATE_TAP_AUTOSTART, _boolean(True)),
        ):
            cf.CFDictionarySetValue(aggregate, cfstr(key), value)

        device_id = ctypes.c_uint32(0)
        status = ca.AudioHardwareCreateAggregateDevice(
            aggregate, ctypes.byref(device_id)
        )
        if status != 0:
            raise TapError(
                f"AudioHardwareCreateAggregateDevice failed ({_osstatus(status)})"
            )
        self._device_id = device_id.value

    def _start_io(self, ca: ctypes.CDLL) -> None:
        self._callback = _IOProc(self._on_audio)
        io_proc = ctypes.c_void_p()
        status = ca.AudioDeviceCreateIOProcID(
            self._device_id, self._callback, None, ctypes.byref(io_proc)
        )
        if status != 0:
            raise TapError(f"AudioDeviceCreateIOProcID failed ({_osstatus(status)})")
        self._io_proc = io_proc
        status = ca.AudioDeviceStart(self._device_id, io_proc)
        if status != 0:
            raise TapError(f"AudioDeviceStart failed ({_osstatus(status)})")

    def _teardown(self) -> None:
        device = self._device_id
        io_proc = self._io_proc
        tap = self._tap_id
        self._device_id = None
        self._io_proc = None
        self._tap_id = None
        self._callback = None
        try:
            ca = _frameworks()[0]
        except OSError:
            return
        if device is not None:
            if io_proc is not None:
                ca.AudioDeviceStop(device, io_proc)
                ca.AudioDeviceDestroyIOProcID(device, io_proc)
            ca.AudioHardwareDestroyAggregateDevice(device)
        if tap is not None:
            ca.AudioHardwareDestroyProcessTap(tap)

    # ----- real-time thread --------------------------------------------- #

    def _on_audio(
        self,
        _device: int,
        _now: int,
        input_data: ctypes.POINTER(_AudioBufferList),
        _input_time: int,
        _output_data: int,
        _output_time: int,
        _client: int,
    ) -> int:
        try:
            chunk = _copy_buffers(input_data)
        except Exception:  # never let an exception escape into the HAL
            return 0
        if not chunk:
            return 0
        with self._cond:
            self._buffers.append(chunk)
            self._buffered += len(chunk)
            while self._buffered > self._limit and self._buffers:
                self._buffered -= len(self._buffers.popleft())
            self._cond.notify_all()
        return 0


def _empty_array() -> int:
    cf = _frameworks()[1]
    return int(cf.CFArrayCreate(None, None, 0, None) or 0)


def _boolean(value: bool) -> int:
    cf = _frameworks()[1]
    name = "kCFBooleanTrue" if value else "kCFBooleanFalse"
    return int(ctypes.c_void_p.in_dll(cf, name).value or 0)


def _get_property(
    ca: ctypes.CDLL,
    obj: int | None,
    selector: int,
    target: Any,
    size: ctypes.c_uint32,
) -> int:
    if obj is None:
        return -1
    address = _PropertyAddress(selector, _SCOPE_GLOBAL, _ELEMENT_MAIN)
    return ca.AudioObjectGetPropertyData(
        obj, ctypes.byref(address), 0, None, ctypes.byref(size), ctypes.byref(target)
    )


def _copy_buffers(input_data: ctypes.POINTER(_AudioBufferList)) -> bytes:
    if not input_data:
        return b""
    block = input_data.contents
    pieces = []
    for index in range(min(int(block.mNumberBuffers), 8)):
        buffer = block.mBuffers[index]
        if buffer.mData and buffer.mDataByteSize:
            pieces.append(ctypes.string_at(buffer.mData, buffer.mDataByteSize))
    return b"".join(pieces)