#!/usr/bin/env python3
"""
G1 AUDIO DEVICES — PulseAudio source/sink enumeration + Hollyland auto-setup.

Pure-Python wrapper around `pactl`. No PulseAudio Python bindings required.
Safe to run on workstation or on the robot.

Capabilities:
  list              — list every PulseAudio source and sink
  find <keyword>    — find a source/sink by substring (case-insensitive)
  auto-hollyland    — detect Hollyland wireless mic, unmute, set default, volume 100%
  quick-test <idx>  — record 2 s from source <idx> via parec and report stats
  mute <idx>        — set-source-mute 1
  unmute <idx>      — set-source-mute 0
  volume <idx> <%>  — set-source-volume <pct>%

See voice_note.txt for why parec is preferred over PyAudio on the robot.

Usage:
  python3 g1_audio_devices.py list
  python3 g1_audio_devices.py find hollyland
  python3 g1_audio_devices.py auto-hollyland
  python3 g1_audio_devices.py quick-test 3
"""

from __future__ import annotations

import argparse
import subprocess
import sys
import tempfile
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable


HOLLYLAND_KEYWORDS = ("hollyland", "wireless_microphone", "shenzhen_hollyland")
ANKER_KEYWORDS = ("anker", "powerconf")


@dataclass
class PulseDevice:
    index: int
    name: str
    description: str
    state: str
    mute: bool | None = None
    volume_pct: int | None = None


def _run(cmd: list[str]) -> tuple[int, str, str]:
    try:
        p = subprocess.run(cmd, capture_output=True, text=True)
    except FileNotFoundError:
        tool = cmd[0]
        print(
            f"error: `{tool}` not on PATH. Install PulseAudio tools "
            f"(e.g. `sudo apt install pulseaudio-utils`) or run this on the robot.",
            file=sys.stderr,
        )
        sys.exit(127)
    return p.returncode, p.stdout, p.stderr


def _list(kind: str) -> list[PulseDevice]:
    """kind = 'sources' or 'sinks'."""
    rc, out, _ = _run(["pactl", "list", "short", kind])
    devices: list[PulseDevice] = []
    if rc != 0:
        return devices
    for line in out.splitlines():
        parts = line.split("\t")
        if len(parts) < 4:
            continue
        idx, name, _driver, _format = parts[0], parts[1], parts[2], parts[3]
        state = parts[4] if len(parts) > 4 else ""
        devices.append(PulseDevice(int(idx), name, "", state))
    # Enrich with mute + volume from `pactl list {sources,sinks}` (long form).
    rc, long_out, _ = _run(["pactl", "list", kind])
    if rc != 0:
        return devices
    blocks = long_out.split("\n\n")
    for block in blocks:
        lines = block.strip().splitlines()
        if not lines:
            continue
        try:
            header_idx = int(lines[0].split("#")[-1].strip())
        except ValueError:
            continue
        mute = None
        volume = None
        desc = ""
        for ln in lines:
            s = ln.strip()
            if s.startswith("Mute:"):
                mute = s.endswith("yes")
            elif s.startswith("Volume:"):
                # "Volume: front-left: 65536 / 100% / 0.00 dB, ..."
                pct_token = None
                for tok in s.split():
                    if tok.endswith("%") and tok[:-1].lstrip("-").isdigit():
                        pct_token = tok
                        break
                if pct_token:
                    try:
                        volume = int(pct_token.rstrip("%"))
                    except ValueError:
                        pass
            elif s.startswith("Description:"):
                desc = s.split(":", 1)[1].strip()
        for d in devices:
            if d.index == header_idx:
                d.mute = mute
                d.volume_pct = volume
                if desc and not d.description:
                    d.description = desc
                break
    return devices


def list_sources() -> list[PulseDevice]:
    return _list("sources")


def list_sinks() -> list[PulseDevice]:
    return _list("sinks")


def find_by_keywords(
    devices: Iterable[PulseDevice], keywords: tuple[str, ...]
) -> PulseDevice | None:
    for d in devices:
        haystack = f"{d.name} {d.description}".lower()
        if any(k in haystack for k in keywords):
            return d
    return None


def pprint_devices(title: str, devices: Iterable[PulseDevice]) -> None:
    print(f"\n{title}")
    print("-" * len(title))
    if not devices:
        print("  (none)")
        return
    for d in devices:
        mute = "muted" if d.mute else "ok"
        vol = f"{d.volume_pct}%" if d.volume_pct is not None else "?"
        print(f"  [{d.index:>2}] {d.name}")
        if d.description:
            print(f"       desc: {d.description}")
        print(f"       state={d.state}  {mute}  vol={vol}")


def set_mute(index: int, mute: bool) -> None:
    _run(["pactl", "set-source-mute", str(index), "1" if mute else "0"])


def set_volume(index: int, pct: int) -> None:
    _run(["pactl", "set-source-volume", str(index), f"{pct}%"])


def set_default_source(name: str) -> None:
    _run(["pactl", "set-default-source", name])


def auto_hollyland() -> int:
    """Follow the voice_note.txt recipe to bring the Hollyland mic online."""
    sources = list_sources()
    mic = find_by_keywords(sources, HOLLYLAND_KEYWORDS)
    if mic is None:
        print("Hollyland mic not found. Plug it in and retry.")
        pprint_devices("Current sources", sources)
        return 1
    print(f"Found: [{mic.index}] {mic.name}")
    set_default_source(mic.name)
    set_mute(mic.index, False)
    set_volume(mic.index, 100)
    print(f"Set as default source, unmuted, volume 100%.")
    return quick_test_capture(mic.index, seconds=2)


def quick_test_capture(index: int, seconds: float = 2.0) -> int:
    """Record a few seconds with parec and report sample statistics."""
    import numpy as np

    with tempfile.NamedTemporaryFile(suffix=".pcm", delete=False) as tmp:
        tmp_path = Path(tmp.name)
    cmd = [
        "parec", "-d", str(index),
        "--format=s16le", "--rate=16000", "--channels=1", "--raw",
    ]
    print(f"Capturing {seconds:.1f}s from source {index} ...")
    proc = subprocess.Popen(cmd, stdout=open(tmp_path, "wb"))
    try:
        proc.wait(timeout=seconds + 2)
    except subprocess.TimeoutExpired:
        proc.terminate()
    finally:
        proc.poll()
    try:
        data = np.fromfile(tmp_path, dtype=np.int16)
    finally:
        tmp_path.unlink(missing_ok=True)

    if data.size == 0:
        print("  No samples captured. Check device index and that parec is installed.")
        return 2
    rms = float((data.astype("float32") ** 2).mean()) ** 0.5
    peak = int(abs(data).max())
    print(f"  samples={data.size}  peak={peak}  rms={rms:.0f}")
    if rms > 50:
        print("  MIC WORKS.")
        return 0
    print("  Signal looks silent — check transmitter is on + mic not muted + volume up.")
    return 3


def main() -> int:
    p = argparse.ArgumentParser(description="G1 audio device helper (pactl wrapper)")
    sub = p.add_subparsers(dest="cmd", required=True)
    sub.add_parser("list", help="List sources and sinks")
    f = sub.add_parser("find", help="Find a device by substring")
    f.add_argument("keyword")
    sub.add_parser("auto-hollyland", help="Auto-detect Hollyland mic and set up")
    q = sub.add_parser("quick-test", help="Record 2s from a source and report")
    q.add_argument("index", type=int)
    q.add_argument("--seconds", type=float, default=2.0)
    m = sub.add_parser("mute", help="Mute a source")
    m.add_argument("index", type=int)
    u = sub.add_parser("unmute", help="Unmute a source")
    u.add_argument("index", type=int)
    v = sub.add_parser("volume", help="Set source volume percent")
    v.add_argument("index", type=int)
    v.add_argument("percent", type=int)
    args = p.parse_args()

    if args.cmd == "list":
        pprint_devices("SOURCES (mics)", list_sources())
        pprint_devices("SINKS (speakers)", list_sinks())
        return 0
    if args.cmd == "find":
        kw = args.keyword.lower()
        src = find_by_keywords(list_sources(), (kw,))
        snk = find_by_keywords(list_sinks(), (kw,))
        if src:
            print(f"source [{src.index}] {src.name}")
        if snk:
            print(f"sink   [{snk.index}] {snk.name}")
        if not src and not snk:
            print(f"No device matches '{args.keyword}'")
            return 1
        return 0
    if args.cmd == "auto-hollyland":
        return auto_hollyland()
    if args.cmd == "quick-test":
        return quick_test_capture(args.index, args.seconds)
    if args.cmd == "mute":
        set_mute(args.index, True)
        return 0
    if args.cmd == "unmute":
        set_mute(args.index, False)
        return 0
    if args.cmd == "volume":
        set_volume(args.index, args.percent)
        return 0
    return 1


if __name__ == "__main__":
    sys.exit(main())