#!/usr/bin/env python3

import subprocess
import re

def parse_wpctl_status():
    devices = []
    current_section = None

    output = subprocess.run(["wpctl", "status"], capture_output=True, text=True).stdout.splitlines()

    for line in output:
        line = line.rstrip()

        if "Sinks:" in line:
            current_section = "sink"
            continue
        elif "Sources:" in line:
            current_section = "source"
            continue
        elif re.match(r"^[^\s]", line):
            current_section = None

        if current_section:
            m = re.search(r"\*\s*|\s*(\d+)\.\s+(.*?)\s+\[vol:", line)
            if m:
                id_match = re.search(r"(\d+)\.", line)
                dev_id = id_match.group(1) if id_match else None
                name_match = re.search(r"\d+\.\s+(.*?)\s+\[vol:", line)
                name = name_match.group(1) if name_match else None

                if dev_id and name:
                    devices.append((current_section, dev_id, name))

    return devices

def main():
    devices = parse_wpctl_status()
    if not devices:
        print("No devices found.")
        return

    # Prepare choices and a mapping from display line to device ID
    choices = []
    id_map = {}

    for dev_type, dev_id, name in devices:
        label = "O" if dev_type == "sink" else "I"
        line = f"{label} {name}"
        choices.append(line)
        id_map[line] = dev_id

    result = subprocess.run(["tofi"], input="\n".join(choices), capture_output=True, text=True)
    choice = result.stdout.strip()
    if not choice:
        return

    selected_id = id_map.get(choice)
    if selected_id:
        subprocess.run(["wpctl", "set-default", selected_id])
        subprocess.run(["notify-send", f"Set default device to {choice}"])

if __name__ == "__main__":
    main()