path: root/linux_audio.py

import os
import select
import shutil
import subprocess
import sys
import threading
import time
from typing import Any, Callable, Dict, List, Optional, Tuple

import numpy as np
import sounddevice as sd

LINUX_VIRTUAL_SINK_MODULE: str = "module-null-sink"
LINUX_VIRTUAL_SOURCE_MODULE: str = "module-remap-source"
LINUX_VIRTUAL_LOOPBACK_MODULE: str = "module-loopback"
LINUX_VIRTUAL_SINK_NAME: str = "auto_live_tl_sink"
LINUX_VIRTUAL_SOURCE_NAME: str = "auto_live_tl_input"
LINUX_VIRTUAL_SINK_DESCRIPTION: str = "auto-live-tl-virtual-sink"
LINUX_VIRTUAL_SOURCE_DESCRIPTION: str = "auto-live-tl-virtual-input"
LINUX_VIRTUAL_DEVICE_WAIT_SECONDS: float = 2.5
LINUX_VIRTUAL_DIRECT_DEVICE_NAME: str = "auto-live-tl-virtual-input (PipeWire direct)"
LINUX_VIRTUAL_DIRECT_CAPTURE_RATE: int = 48000
LINUX_DIRECT_CAPTURE_IDLE_TIMEOUT_SECONDS: float = 3.0

_virtual_sink_module_id: Optional[int] = None
_virtual_source_module_id: Optional[int] = None
_virtual_loopback_module_id: Optional[int] = None
_virtual_sink_created_by_app: bool = False
_virtual_source_created_by_app: bool = False
_virtual_loopback_created_by_app: bool = False


def is_linux() -> bool:
    return sys.platform.startswith("linux")


def is_virtual_input_name(device_name: str) -> bool:
    lowered = device_name.lower()
    return (
        LINUX_VIRTUAL_SOURCE_NAME.lower() in lowered
        or LINUX_VIRTUAL_SOURCE_DESCRIPTION.lower() in lowered
        or LINUX_VIRTUAL_SINK_NAME.lower() in lowered
        or LINUX_VIRTUAL_SINK_DESCRIPTION.lower() in lowered
    )


def _run_pactl(args: List[str]) -> subprocess.CompletedProcess[str]:
    return subprocess.run(
        ["pactl", *args],
        check=False,
        capture_output=True,
        text=True,
    )


def _find_existing_module_id(module_name: str, arg_match: str) -> Optional[int]:
    result = _run_pactl(["list", "short", "modules"])
    if result.returncode != 0:
        return None

    for line in result.stdout.splitlines():
        parts = line.split("\t")
        if len(parts) < 3:
            continue
        module_id_raw, listed_module_name, module_args = parts[0], parts[1], parts[2]
        if listed_module_name != module_name:
            continue
        if arg_match not in module_args:
            continue
        try:
            return int(module_id_raw)
        except ValueError:
            continue
    return None


def _wait_for_virtual_input_device() -> bool:
    deadline = time.monotonic() + LINUX_VIRTUAL_DEVICE_WAIT_SECONDS
    while time.monotonic() < deadline:
        try:
            devices = sd.query_devices()
        except Exception:
            time.sleep(0.1)
            continue
        for dev in devices:
            if dev.get("max_input_channels", 0) > 0 and is_virtual_input_name(str(dev.get("name", ""))):
                return True
        time.sleep(0.1)
    return False


def _get_default_sink_name() -> str:
    result = _run_pactl(["get-default-sink"])
    if result.returncode == 0:
        sink_name = (result.stdout or "").strip()
        if sink_name:
            return sink_name
    return "@DEFAULT_SINK@"


def ensure_virtual_input_sink() -> None:
    global _virtual_sink_module_id, _virtual_source_module_id, _virtual_loopback_module_id
    global _virtual_sink_created_by_app, _virtual_source_created_by_app, _virtual_loopback_created_by_app

    if not is_linux():
        return

    if shutil.which("pactl") is None:
        print("[audio] 'pactl' not found, skipping PipeWire virtual device creation.")
        return

    existing_sink_id = _find_existing_module_id(
        LINUX_VIRTUAL_SINK_MODULE,
        f"sink_name={LINUX_VIRTUAL_SINK_NAME}",
    )
    if existing_sink_id is not None:
        _virtual_sink_module_id = existing_sink_id
        _virtual_sink_created_by_app = False
        print(f"[audio] Reusing existing virtual sink module id={existing_sink_id}.")
    else:
        sink_result = _run_pactl(
            [
                "load-module",
                LINUX_VIRTUAL_SINK_MODULE,
                f"sink_name={LINUX_VIRTUAL_SINK_NAME}",
                f"sink_properties=device.description={LINUX_VIRTUAL_SINK_DESCRIPTION}",
            ]
        )
        if sink_result.returncode != 0:
            stderr = (sink_result.stderr or "").strip()
            print(f"[audio] Failed to create PipeWire virtual sink: {stderr or 'unknown error'}")
            return

        sink_id_text = (sink_result.stdout or "").strip()
        try:
            _virtual_sink_module_id = int(sink_id_text)
            _virtual_sink_created_by_app = True
            print(
                "[audio] Created PipeWire virtual sink "
                f"'{LINUX_VIRTUAL_SINK_DESCRIPTION}' (module id={_virtual_sink_module_id})."
            )
        except ValueError:
            _virtual_sink_module_id = None
            _virtual_sink_created_by_app = True
            print("[audio] Created PipeWire virtual sink, but could not parse module id.")

    existing_source_id = _find_existing_module_id(
        LINUX_VIRTUAL_SOURCE_MODULE,
        f"source_name={LINUX_VIRTUAL_SOURCE_NAME}",
    )
    if existing_source_id is not None:
        _virtual_source_module_id = existing_source_id
        _virtual_source_created_by_app = False
        print(f"[audio] Reusing existing virtual input module id={existing_source_id}.")
    else:
        source_result = _run_pactl(
            [
                "load-module",
                LINUX_VIRTUAL_SOURCE_MODULE,
                f"master={LINUX_VIRTUAL_SINK_NAME}.monitor",
                f"source_name={LINUX_VIRTUAL_SOURCE_NAME}",
                f"source_properties=device.description={LINUX_VIRTUAL_SOURCE_DESCRIPTION}",
            ]
        )
        if source_result.returncode != 0:
            stderr = (source_result.stderr or "").strip()
            print(f"[audio] Failed to create virtual input source: {stderr or 'unknown error'}")
            return

        source_id_text = (source_result.stdout or "").strip()
        try:
            _virtual_source_module_id = int(source_id_text)
            _virtual_source_created_by_app = True
            print(
                "[audio] Created virtual input source "
                f"'{LINUX_VIRTUAL_SOURCE_DESCRIPTION}' (module id={_virtual_source_module_id})."
            )
        except ValueError:
            _virtual_source_module_id = None
            _virtual_source_created_by_app = True
            print("[audio] Created virtual input source, but could not parse module id.")

    existing_loopback_id = _find_existing_module_id(
        LINUX_VIRTUAL_LOOPBACK_MODULE,
        f"source={LINUX_VIRTUAL_SINK_NAME}.monitor",
    )
    if existing_loopback_id is not None:
        _virtual_loopback_module_id = existing_loopback_id
        _virtual_loopback_created_by_app = False
        print(f"[audio] Reusing existing headphone loopback module id={existing_loopback_id}.")
    else:
        default_sink = _get_default_sink_name()
        loopback_result = _run_pactl(
            [
                "load-module",
                LINUX_VIRTUAL_LOOPBACK_MODULE,
                f"source={LINUX_VIRTUAL_SINK_NAME}.monitor",
                f"sink={default_sink}",
                "latency_msec=60",
            ]
        )
        if loopback_result.returncode != 0:
            stderr = (loopback_result.stderr or "").strip()
            print(f"[audio] Failed to create headphone loopback: {stderr or 'unknown error'}")
        else:
            loopback_id_text = (loopback_result.stdout or "").strip()
            try:
                _virtual_loopback_module_id = int(loopback_id_text)
                _virtual_loopback_created_by_app = True
                print(
                    "[audio] Mirroring virtual sink to default output "
                    f"(module id={_virtual_loopback_module_id}, sink={default_sink})."
                )
            except ValueError:
                _virtual_loopback_module_id = None
                _virtual_loopback_created_by_app = True
                print("[audio] Headphone loopback created, but could not parse module id.")

    if not _wait_for_virtual_input_device():
        print("[audio] Virtual input device was not detected by PortAudio yet.")


def teardown_virtual_input_sink() -> None:
    global _virtual_sink_module_id, _virtual_source_module_id, _virtual_loopback_module_id
    global _virtual_sink_created_by_app, _virtual_source_created_by_app, _virtual_loopback_created_by_app

    if not is_linux() or shutil.which("pactl") is None:
        return

    if _virtual_loopback_created_by_app and _virtual_loopback_module_id is not None:
        result = _run_pactl(["unload-module", str(_virtual_loopback_module_id)])
        if result.returncode != 0:
            stderr = (result.stderr or "").strip()
            print(f"[audio] Failed to unload headphone loopback module {_virtual_loopback_module_id}: {stderr or 'unknown error'}")
        else:
            print(f"[audio] Unloaded headphone loopback module {_virtual_loopback_module_id}.")

    if _virtual_source_created_by_app and _virtual_source_module_id is not None:
        result = _run_pactl(["unload-module", str(_virtual_source_module_id)])
        if result.returncode != 0:
            stderr = (result.stderr or "").strip()
            print(f"[audio] Failed to unload virtual input module {_virtual_source_module_id}: {stderr or 'unknown error'}")
        else:
            print(f"[audio] Unloaded virtual input module {_virtual_source_module_id}.")

    if _virtual_sink_created_by_app and _virtual_sink_module_id is not None:
        result = _run_pactl(["unload-module", str(_virtual_sink_module_id)])
        if result.returncode != 0:
            stderr = (result.stderr or "").strip()
            print(f"[audio] Failed to unload virtual sink module {_virtual_sink_module_id}: {stderr or 'unknown error'}")
        else:
            print(f"[audio] Unloaded virtual sink module {_virtual_sink_module_id}.")

    _virtual_loopback_module_id = None
    _virtual_source_module_id = None
    _virtual_sink_module_id = None
    _virtual_loopback_created_by_app = False
    _virtual_source_created_by_app = False
    _virtual_sink_created_by_app = False


def append_pipewire_direct_device_option(
    input_devices: List[Tuple[int, Dict[str, Any]]],
    has_virtual_input: bool,
) -> None:
    if not has_virtual_input:
        print(
            "[audio] Virtual input device not visible in sounddevice device list. "
            "Adding PipeWire direct capture option instead."
        )

    input_devices.append(
        (
            -1,
            {
                "name": LINUX_VIRTUAL_DIRECT_DEVICE_NAME,
                "max_input_channels": 1,
                "default_samplerate": float(LINUX_VIRTUAL_DIRECT_CAPTURE_RATE),
            },
        )
    )


def _build_pipewire_direct_commands(sample_rate: int) -> List[List[str]]:
    commands: List[List[str]] = []

    if shutil.which("ffmpeg") is not None:
        for source_name in [LINUX_VIRTUAL_SOURCE_NAME, f"{LINUX_VIRTUAL_SINK_NAME}.monitor"]:
            commands.append(
                [
                    "ffmpeg",
                    "-nostdin",
                    "-hide_banner",
                    "-loglevel",
                    "error",
                    "-f",
                    "pulse",
                    "-i",
                    source_name,
                    "-ac",
                    "1",
                    "-ar",
                    str(sample_rate),
                    "-f",
                    "f32le",
                    "-",
                ]
            )

    return commands


def has_pipewire_direct_backend(sample_rate: int) -> bool:
    return bool(_build_pipewire_direct_commands(sample_rate))


def pipewire_direct_capture_worker(
    stop_event: threading.Event,
    sample_rate: int,
    on_chunk: Callable[[np.ndarray], None],
    add_runtime_log: Callable[[str, str], None],
) -> None:
    try:
        _run_pipewire_direct_capture(stop_event, sample_rate, on_chunk)
    except Exception as exc:
        message = f"PipeWire direct capture failed: {exc}"
        print(f"[audio] {message}")
        add_runtime_log("AUDIO", message)


def _run_pipewire_direct_capture(
    stop_event: threading.Event,
    sample_rate: int,
    on_chunk: Callable[[np.ndarray], None],
) -> None:
    commands = _build_pipewire_direct_commands(sample_rate)
    if not commands:
        raise RuntimeError("PipeWire direct capture needs 'ffmpeg' installed.")

    bytes_per_chunk = int(sample_rate * 0.5) * 4

    for command in commands:
        process: Optional[subprocess.Popen[bytes]] = None
        try:
            process = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            time.sleep(0.2)
            if process.poll() is not None:
                continue

            assert process.stdout is not None
            print(f"[audio] PipeWire direct capture backend: {' '.join(command[:2])}")

            silent_windows = 0
            while not stop_event.is_set():
                ready, _, _ = select.select([process.stdout], [], [], 1.0)
                if not ready:
                    silent_windows += 1
                    if silent_windows * 1.0 >= LINUX_DIRECT_CAPTURE_IDLE_TIMEOUT_SECONDS:
                        raise RuntimeError("backend produced no audio frames")
                    continue

                raw = os.read(process.stdout.fileno(), bytes_per_chunk)
                if not raw:
                    raise RuntimeError("backend stream ended")

                silent_windows = 0
                chunk = np.frombuffer(raw, dtype=np.float32).copy()
                if chunk.size == 0:
                    continue
                on_chunk(chunk)

            return
        except Exception as exc:
            backend_name = " ".join(command[:2])
            print(f"[audio] Direct capture backend failed ({backend_name}): {exc}")
            continue
        finally:
            if process is not None and process.poll() is None:
                process.terminate()
                try:
                    process.wait(timeout=1.5)
                except subprocess.TimeoutExpired:
                    process.kill()

    raise RuntimeError("All PipeWire direct capture backends failed.")