#!/usr/bin/env python3
import json
import signal
import sys
import threading
import time
from pathlib import Path

import jack
import yaml
import paho.mqtt.client as mqtt


CONFIG_FILE = Path(__file__).resolve().with_name("config.yaml")


class JackRouter:
    def __init__(self, config_file=CONFIG_FILE, jack_client=None):
        with open(config_file, "r", encoding="utf-8") as f:
            self.cfg = yaml.safe_load(f)

        self.client = jack_client or jack.Client("AudioRouter")
        self.lock = threading.RLock()
        self.mode = self.cfg.get("router", {}).get("startup_mode", "mute_all")
        self.base_topic = self.cfg["mqtt"]["base_topic"].rstrip("/")
        self._reported_failures = set()
        self.exclusive_sources = set(
            self.cfg.get("router", {}).get("exclusive_sources", ["ytplayer"])
        )
        self._validate_config()

    def _validate_config(self):
        routes = self.cfg["routes"]
        ports = self.cfg["ports"]

        for route_name, route in routes.items():
            if route["source"] not in ports or route["output"] not in ports:
                raise ValueError(f"Route {route_name!r} references an unknown port group")

        for mode, preset in self.cfg["presets"].items():
            unknown_routes = set(preset) - set(routes)
            if unknown_routes:
                raise ValueError(
                    f"Preset {mode!r} references unknown routes: "
                    f"{sorted(unknown_routes)!r}"
                )

            enabled_outputs = {}
            for route_name, enabled in preset.items():
                if not enabled:
                    continue
                route = routes[route_name]
                enabled_outputs.setdefault(route["source"], set()).add(route["output"])

            for source in self.exclusive_sources:
                outputs = enabled_outputs.get(source, set())
                if len(outputs) > 1:
                    raise ValueError(
                        f"Preset {mode!r} routes exclusive source {source!r} "
                        f"to multiple outputs: {sorted(outputs)!r}"
                    )

    def _log_route_failure(self, operation, src, dst, error):
        key = (operation, src, dst)
        if key not in self._reported_failures:
            print(f"{operation.capitalize()} failed: {src} -> {dst}: {error}", flush=True)
            self._reported_failures.add(key)

    def _clear_route_failure(self, operation, src, dst):
        self._reported_failures.discard((operation, src, dst))

    def connect(self, src, dst):
        try:
            connections = self.client.get_all_connections(dst)
            connected_names = [port.name for port in connections]

            if src not in connected_names:
                self.client.connect(src, dst)
                print(f"Connected: {src} -> {dst}", flush=True)
            self._clear_route_failure("connect", src, dst)
            return True
        except jack.JackError as e:
            self._log_route_failure("connect", src, dst, e)
            return False

    def connection_state(self, src, dst):
        """Return True/False for a link, or None when it cannot be inspected."""
        try:
            connections = self.client.get_all_connections(dst)
            connected_names = [port.name for port in connections]
            self._clear_route_failure("inspect", src, dst)
            return src in connected_names
        except jack.JackError as error:
            self._log_route_failure("inspect", src, dst, error)
            return None

    def disconnect(self, src, dst):
        try:
            self.client.disconnect(src, dst)
            print(f"Disconnected: {src} -> {dst}", flush=True)
        except jack.JackError:
            pass

        state = self.connection_state(src, dst)
        if state is False:
            self._clear_route_failure("disconnect", src, dst)
            return True
        if state is True:
            self._log_route_failure(
                "disconnect",
                src,
                dst,
                "link is still present after disconnect",
            )
        return False

    def route_pairs(self, route_name):
        route = self.cfg["routes"][route_name]
        source = self.cfg["ports"][route["source"]]
        output = self.cfg["ports"][route["output"]]
        return [
            (source["left"], output["left"]),
            (source["right"], output["right"]),
        ]

    def all_route_pairs(self):
        pairs = []
        for route_name in self.cfg["routes"]:
            pairs.extend(self.route_pairs(route_name))
        return pairs

    def clear_routes(self):
        with self.lock:
            for src, dst in self.all_route_pairs():
                self.disconnect(src, dst)

    def routes_for_mode(self, mode):
        route_names = self.route_names_for_mode(mode)

        if route_names is None:
            return None

        resolved_routes = []
        for route_name in route_names:
            resolved_routes.extend(self.route_pairs(route_name))

        return resolved_routes

    def route_names_for_mode(self, mode):
        preset = self.cfg["presets"].get(mode)
        if preset is None:
            return None
        return [route_name for route_name, enabled in preset.items() if enabled]

    def exclusive_route_is_safe(self, desired_route_name):
        desired_route = self.cfg["routes"][desired_route_name]
        source_name = desired_route["source"]
        if source_name not in self.exclusive_sources:
            return True

        desired_output = desired_route["output"]
        safe = True
        for route_name, route in self.cfg["routes"].items():
            if route["source"] != source_name or route["output"] == desired_output:
                continue

            for src, dst in self.route_pairs(route_name):
                state = self.connection_state(src, dst)
                if state is True:
                    self.disconnect(src, dst)
                    state = self.connection_state(src, dst)
                if state is not False:
                    self._log_route_failure(
                        "exclusive-check",
                        src,
                        dst,
                        f"cannot safely route {source_name!r} to {desired_output!r}",
                    )
                    safe = False
        return safe

    def connect_routes(self, route_names):
        for route_name in route_names:
            if not self.exclusive_route_is_safe(route_name):
                continue
            for src, dst in self.route_pairs(route_name):
                self.connect(src, dst)

    def set_mode(self, mode):
        route_names = self.route_names_for_mode(mode)
        if route_names is None:
            print(f"Unknown mode: {mode}", flush=True)
            return False

        with self.lock:
            if mode == self.mode:
                print(f"Reconciling mode: {mode}", flush=True)
                self.reconcile()
                return True

            print(f"Setting mode: {mode}", flush=True)
            # Break before make: remove every managed route before applying the
            # next preset, so an exclusive source can never bridge two outputs.
            self.clear_routes()
            self.mode = mode
            self.connect_routes(route_names)

        return True

    def reconcile(self):
        """Restore the selected preset after PipeWire ports appear or restart."""
        with self.lock:
            route_names = self.route_names_for_mode(self.mode) or []
            desired = set(self.routes_for_mode(self.mode) or [])

            for src, dst in self.all_route_pairs():
                if (src, dst) not in desired:
                    self.disconnect(src, dst)

            self.connect_routes(route_names)

    def state(self):
        with self.lock:
            return {
                "mode": self.mode,
                "running": True,
                "timestamp": int(time.time()),
            }


def main():
    router = JackRouter()
    shutdown_event = threading.Event()
    mqtt_connected_event = threading.Event()
    startup_restored_event = threading.Event()
    restore_lock = threading.Lock()
    restore_state = {"deadline": None}

    mqtt_cfg = router.cfg["mqtt"]
    client = mqtt.Client(mqtt.CallbackAPIVersion.VERSION2)
    mode_set_topic = f"{router.base_topic}/mode/set"
    mode_state_topic = f"{router.base_topic}/mode/state"

    def publish_state():
        snapshot = router.state()
        client.publish(f"{router.base_topic}/state", json.dumps(snapshot), retain=True)
        client.publish(mode_state_topic, snapshot["mode"], retain=True)

    def finish_startup_restore():
        with restore_lock:
            if startup_restored_event.is_set():
                return False
            startup_restored_event.set()
            restore_state["deadline"] = None

        client.unsubscribe(mode_state_topic)
        return True

    def on_connect(client, userdata, flags, reason_code, properties=None):
        try:
            if reason_code != 0:
                print(f"MQTT connection failed: {reason_code}", file=sys.stderr, flush=True)
                return

            print(f"MQTT connected: {reason_code}", flush=True)
            mqtt_connected_event.set()
            client.subscribe(mode_set_topic)

            if startup_restored_event.is_set():
                publish_state()
            else:
                # mode/set is intentionally not retained on this installation,
                # while mode/state is. Read mode/state once to restore the last
                # applied preset, then unsubscribe so it remains output-only.
                client.subscribe(mode_state_topic)
                with restore_lock:
                    restore_state["deadline"] = time.monotonic() + 2.0
        except Exception as error:
            print(f"MQTT on_connect error: {error}", file=sys.stderr, flush=True)

    def on_connect_fail(client, userdata):
        print(
            f"MQTT connection to {mqtt_cfg['host']}:{mqtt_cfg['port']} failed; "
            "will retry",
            file=sys.stderr,
            flush=True,
        )

    def on_disconnect(client, userdata, disconnect_flags, reason_code, properties=None):
        mqtt_connected_event.clear()
        with restore_lock:
            if not startup_restored_event.is_set():
                restore_state["deadline"] = None
        if shutdown_event.is_set():
            print(f"MQTT disconnected: {reason_code}", flush=True)
        else:
            print(f"MQTT disconnected: {reason_code}; will retry", file=sys.stderr, flush=True)

    def on_message(client, userdata, msg):
        try:
            if shutdown_event.is_set():
                return
            mode = msg.payload.decode("utf-8", errors="replace").strip()
            print(f"MQTT received: {msg.topic} = {mode}", flush=True)

            if msg.topic == mode_state_topic:
                if (
                    getattr(msg, "retain", False)
                    and not startup_restored_event.is_set()
                    and router.set_mode(mode)
                ):
                    finish_startup_restore()
                    publish_state()
                return

            if msg.topic == mode_set_topic and router.set_mode(mode):
                finish_startup_restore()
                publish_state()
        except Exception as error:
            print(f"MQTT on_message error: {error}", file=sys.stderr, flush=True)

    def shutdown(*args):
        print("Stopping router...", flush=True)
        shutdown_event.set()

    signal.signal(signal.SIGINT, shutdown)
    signal.signal(signal.SIGTERM, shutdown)

    client.on_connect = on_connect
    client.on_connect_fail = on_connect_fail
    client.on_disconnect = on_disconnect
    client.on_message = on_message
    client.suppress_exceptions = True
    client.reconnect_delay_set(min_delay=1, max_delay=30)

    jack_shutdown_state = {"stopped": False, "reason": None}

    def on_jack_shutdown(status, reason):
        # JACK invokes this from its process thread. Only update plain state;
        # the main thread will exit and systemd will construct a fresh client.
        jack_shutdown_state["stopped"] = True
        jack_shutdown_state["reason"] = reason

    router.client.set_shutdown_callback(on_jack_shutdown)

    with router.client:
        # Clear any stale managed links before network activity. Desired mode
        # will be restored from the retained mode/state publication.
        router.set_mode(router.mode)
        client.connect_async(mqtt_cfg["host"], int(mqtt_cfg["port"]), 60)
        loop_rc = client.loop_start()
        if loop_rc != mqtt.MQTT_ERR_SUCCESS:
            raise RuntimeError(f"Failed to start MQTT network loop: {loop_rc}")

        interval = max(
            0.2,
            float(router.cfg.get("router", {}).get("reconcile_interval", 1.0)),
        )
        try:
            while not shutdown_event.wait(interval):
                if jack_shutdown_state["stopped"]:
                    raise RuntimeError(
                        f"JACK/PipeWire stopped: {jack_shutdown_state['reason']}"
                    )
                with restore_lock:
                    restore_deadline = restore_state["deadline"]
                if (
                    mqtt_connected_event.is_set()
                    and not startup_restored_event.is_set()
                    and restore_deadline is not None
                    and time.monotonic() >= restore_deadline
                ):
                    # No retained mode was available. Remain safely muted and
                    # publish that truthful state until a new command arrives.
                    if finish_startup_restore():
                        publish_state()
                router.reconcile()
        finally:
            shutdown_event.set()
            client.loop_stop()
            # Disconnect live links without overwriting the retained desired
            # mode. The next process can then restore that preset safely.
            if not jack_shutdown_state["stopped"]:
                router.clear_routes()


if __name__ == "__main__":
    main()