#!/usr/bin/python3


from typing import Deque, Dict

import collections
import logging
import os

import pulsectl  # type: ignore


logger = logging.getLogger(os.path.basename(__file__))
warn = logger.warn
info = logger.info
debug = logger.debug


class AutoSurrounder(object):
    def __init__(self) -> None:
        self.event_queue: Deque[pulsectl.pulsectl.PulseEventInfo] = collections.deque()
        self.pulse: pulsectl.Pulse = None
        self.multichannel_sink_inputs: Dict[int, set[int]] = collections.defaultdict(
            set
        )
        self.sink_cache: Dict[int, pulsectl.pulsectl.PulseSinkInfo] = {}
        self.card_cache: Dict[int, pulsectl.pulsectl.PulseCardInfo] = {}
        self.saved_card_profiles: Dict[int, pulsectl.pulsectl.PulseCardProfileInfo] = {}

    def get_sink(self, sink_index: int) -> pulsectl.pulsectl.PulseSinkInfo | None:
        if sink_index not in self.sink_cache:
            for x in self.pulse.sink_list():
                self.sink_cache[x.index] = x
        return self.sink_cache.get(sink_index)

    def get_card(self, card_index: int) -> pulsectl.pulsectl.PulseCardInfo | None:
        if card_index not in self.card_cache:
            for x in self.pulse.card_list():
                self.card_cache[x.index] = x
        return self.card_cache.get(card_index)

    def get_sink_input(
        self, sink_input_index: int
    ) -> pulsectl.pulsectl.PulseSinkInputInfo | None:
        res = [x for x in self.pulse.sink_input_list() if x.index == sink_input_index]
        return res[0] if res else None

    def process_event(self, ev: pulsectl.pulsectl.PulseEventInfo) -> bool:
        def is_multichannel(sink_input: pulsectl.pulsectl.PulseSinkInfo) -> bool:
            return bool(sink_input.channel_count > 2)

        def new(ev: pulsectl.pulsectl.PulseEventInfo) -> bool:
            return bool(ev.t == pulsectl.PulseEventTypeEnum.new)

        def changed(ev: pulsectl.pulsectl.PulseEventInfo) -> bool:
            return bool(ev.t == pulsectl.PulseEventTypeEnum.change)

        def removed(ev: pulsectl.pulsectl.PulseEventInfo) -> bool:
            return bool(ev.t == pulsectl.PulseEventTypeEnum.remove)

        def is_sink_input(ev: pulsectl.pulsectl.PulseEventInfo) -> bool:
            return bool(ev.facility == pulsectl.PulseEventFacilityEnum.sink_input)

        def is_sink(ev: pulsectl.pulsectl.PulseEventInfo) -> bool:
            return bool(ev.facility == pulsectl.PulseEventFacilityEnum.sink)

        def is_card(ev: pulsectl.pulsectl.PulseEventInfo) -> bool:
            return bool(ev.facility == pulsectl.PulseEventFacilityEnum.card)

        has_multichannel_appeared = False

        if is_sink(ev):
            if changed(ev):
                if ev.index in self.sink_cache:
                    del self.sink_cache[ev.index]
            if removed(ev):
                if ev.index in self.sink_cache:
                    del self.sink_cache[ev.index]
                if ev.index in self.multichannel_sink_inputs:
                    del self.multichannel_sink_inputs[ev.index]

        elif is_card(ev):
            if changed(ev):
                if ev.index in self.card_cache:
                    del self.card_cache[ev.index]
            if removed(ev):
                if ev.index in self.card_cache:
                    del self.card_cache[ev.index]
                if ev.index in self.saved_card_profiles:
                    del self.saved_card_profiles[ev.index]

        elif is_sink_input(ev):
            # Clear the index from all input slots preemptively.
            if new(ev) or changed(ev):
                sink_input = self.get_sink_input(ev.index)
                if sink_input:
                    sink = self.get_sink(sink_input.sink)
                    if sink:
                        if is_multichannel(sink_input):
                            # (Re-)add the index to the corresponding card.
                            self.multichannel_sink_inputs[sink_input.sink].add(ev.index)
                            has_multichannel_appeared = True
                    else:
                        # print(
                        #    "Sink #%s has no sink %s, weird"
                        #    % (ev.index, sink_input.sink),
                        #    file=sys.stderr,
                        # )
                        self.evict_sink_input(ev.index)
                else:
                    pass
                    # print("Sink input #%s has vanished, weird" % ev.index,
                    #       file=sys.stderr)
            elif removed(ev):
                # print("Sink input #%s is gone" % ev.index, file=sys.stderr)
                self.evict_sink_input(ev.index)

        return has_multichannel_appeared

    def evict_sink_input(self, sink_input_index: int) -> None:
        for unused_sink, sink_input_indexes in self.multichannel_sink_inputs.items():
            if sink_input_index in sink_input_indexes:
                sink_input_indexes.remove(sink_input_index)

    def apply_card_profile_changes(self) -> None:
        cards_multichannel: Dict[
            pulsectl.PulseCardInfo, bool
        ] = collections.defaultdict(bool)
        for sink_index, mc_input_indexes in self.multichannel_sink_inputs.items():
            sink = self.get_sink(sink_index)
            if sink is None:
                # print("Sink %s is none, weird" % sink_index)
                continue
            card = self.get_card(sink.card)
            if card is None:
                # print("Card %s is none, weird" % sink.card)
                continue
            debug(
                "Sink %s (%s) of card %s has these multichannel sink inputs: %s",
                sink,
                sink_index,
                card,
                mc_input_indexes,
            )
            cards_multichannel[card] = (
                bool(mc_input_indexes) or cards_multichannel[card]
            )

        for card, needs_multichannel in cards_multichannel.items():
            profiles = [x for x in card.profile_list if x.available]
            surround_profiles = [x for x in profiles if "surround" in x.name]
            if len(profiles) < 2:
                # This card does not have a profile to switch from/to.
                # print("No profiles to switch from/to")
                continue
            if len(surround_profiles) < 1:
                # print("No surround profiles")
                # This card does not have surround profiles available.
                continue

            is_currently_multichannel = "surround" in card.profile_active.name
            if not is_currently_multichannel and needs_multichannel:
                # Switch TO multichannel.
                old_profile = card.profile_active
                new_profile = surround_profiles[0]
                if card.index not in self.saved_card_profiles:
                    self.saved_card_profiles[card.index] = old_profile
                info(
                    "Switching %s from %s to %s"
                    % (card.name, old_profile.description, new_profile.description),
                )
                self.pulse.card_profile_set(card, new_profile.name)
            elif is_currently_multichannel and not needs_multichannel:
                # Switch FROM multichannel to prior profile.
                if card.index not in self.saved_card_profiles:
                    warn(
                        "No known old profile to restore the %s to." % (card.name),
                    )
                else:
                    old_profile = card.profile_active
                    new_profile = self.saved_card_profiles[card.index]
                    del self.saved_card_profiles[card.index]
                    info(
                        "Switching %s from %s to %s"
                        % (card.name, old_profile.description, new_profile.description),
                    )
                    self.pulse.card_profile_set(card, new_profile.name)

    def run(self) -> None:
        def queue_event(ev: pulsectl.pulsectl.PulseEventInfo) -> None:
            self.event_queue.append(ev)
            raise pulsectl.PulseLoopStop()

        with pulsectl.Pulse("pautosurround") as pulse:
            self.pulse = pulse
            pulse.event_mask_set(
                pulsectl.PulseEventMaskEnum.card,
                pulsectl.PulseEventMaskEnum.sink,
                pulsectl.PulseEventMaskEnum.sink_input,
            )
            pulse.event_callback_set(lambda ev: queue_event(ev))
            short_timeout = 2
            long_timeout = 30
            timeout = long_timeout
            while True:
                multichannel_appeared = False
                timed_out = True
                pulse.event_listen(timeout=timeout)
                while len(self.event_queue):
                    timed_out = False
                    ev = self.event_queue.popleft()
                    multichannel_appeared = (
                        self.process_event(ev) or multichannel_appeared
                    )
                if multichannel_appeared or timed_out:
                    debug(
                        "Multichannel appeared: %s, timed out: %s"
                        % (multichannel_appeared, timed_out),
                    )
                    self.apply_card_profile_changes()
                if multichannel_appeared:
                    # Ensure the timeout is shortened to restore cards to
                    # stereo faster.
                    timeout = short_timeout
                elif timed_out:
                    # Lengthen the future timeout to avoid spending
                    # CPU resources on processing events.
                    timeout = long_timeout


def main() -> None:
    logging.basicConfig(level=logging.DEBUG if os.getenv("DEBUG") else logging.INFO)
    a = AutoSurrounder()
    a.run()


if __name__ == "__main__":
    main()
