Musubi/Packages/MusubiKit/Sources/AudioMatrix/AudioMatrix.swift

import CAudioRing
import Foundation
import Observation
import OSLog

/// Routes audio sources (NDI, input devices) to single channels of output devices.
///
/// NDI audio is pushed in by the caller through the `AudioFanout` returned from `fanout(forNDI:)`;
/// input devices are captured here. Input-device routes are gated by the talk key.
@MainActor @Observable
public final class AudioMatrix {
    public enum SourceRef: Hashable, Sendable {
        case ndi(name: String)
        case input(deviceUID: String)
    }

    public struct Route: Identifiable, Equatable, Sendable {
        public let id: UUID
        public let source: SourceRef
        public let deviceUID: String
        public let channel: Int
        public var selection: ChannelSelection
        public var gainDB: Float
        public var muted: Bool
    }

    public private(set) var devices: [AudioDevice] = []
    public private(set) var routes: [Route] = []
    public private(set) var errors: [String: String] = [:] // device UID → message
    public var talkActive = false { didSet { applyGates() } }

    /// Fanouts for NDI sources; the app feeds them from NDI receivers.
    public private(set) var ndiFanouts: [String: AudioFanout] = [:]

    private var renderers: [UUID: RouteRenderer] = [:]
    /// Removed routes stay alive briefly so the render thread never releases (frees) them.
    private var retired: [RouteRenderer] = []
    private var outputEngines: [String: OutputEngine] = [:]
    private var inputEngines: [String: InputEngine] = [:]
    private var deviceObserver: AnyObject?
    private let log = Logger(subsystem: "studio.zielke.musubi", category: "audio")

    public init() {
        devices = AudioDevices.all()
        deviceObserver = AudioDevices.observeChanges { [weak self] in self?.devicesChanged() }
    }

    public var outputDevices: [AudioDevice] { devices.filter { $0.outputChannels > 0 } }
    public var inputDevices: [AudioDevice] { devices.filter { $0.inputChannels > 0 } }

    @discardableResult
    public func addRoute(_ source: SourceRef, deviceUID: String, channel: Int, selection: ChannelSelection = .mono) -> UUID {
        let id = UUID()
        let latency: Double = if case .input = source { 0.012 } else { 0.08 }
        let feed = AudioFeed()
        let renderer = RouteRenderer(id: id, feed: feed, outputChannel: channel, selection: selection, targetLatency: latency)
        renderers[id] = renderer
        routes.append(Route(id: id, source: source, deviceUID: deviceUID, channel: channel, selection: selection, gainDB: 0, muted: false))

        switch source {
        case .ndi(let name):
            let fanout = ndiFanouts[name] ?? AudioFanout()
            ndiFanouts[name] = fanout
            fanout.add(feed)
        case .input(let uid):
            inputEngine(for: uid)?.fanout.add(feed)
        }
        applyGates()
        rebuildOutputs()
        log.info("Route \(String(describing: source), privacy: .public) → \(deviceUID, privacy: .public) ch \(channel + 1)")
        return id
    }

    public func removeRoute(_ id: UUID) {
        guard let index = routes.firstIndex(where: { $0.id == id }), let renderer = renderers.removeValue(forKey: id) else { return }
        let route = routes.remove(at: index)
        retire(renderer)
        switch route.source {
        case .ndi(let name):
            ndiFanouts[name]?.remove(renderer.feed)
            if ndiFanouts[name]?.isEmpty == true { ndiFanouts[name] = nil }
        case .input(let uid):
            inputEngines[uid]?.fanout.remove(renderer.feed)
            if !routes.contains(where: { $0.source == route.source }) {
                inputEngines.removeValue(forKey: uid)?.stop()
            }
        }
        rebuildOutputs()
    }

    public func setGain(_ id: UUID, db: Float) {
        update(id) { $0.gainDB = db }
        renderers[id]?.gain = Level.gain(fromDB: db)
    }

    public func setMuted(_ id: UUID, _ muted: Bool) {
        update(id) { $0.muted = muted }
        renderers[id]?.muted = muted
    }

    public func setSelection(_ id: UUID, _ selection: ChannelSelection) {
        update(id) { $0.selection = selection }
        renderers[id]?.selection = selection
    }

    /// Peak level of a route since the last call, in dBFS.
    public func takePeak(_ id: UUID) -> Float {
        guard let renderer = renderers[id] else { return Level.floorDB }
        let peak = renderer.peak
        renderer.peak = 0
        return Level.db(fromGain: peak)
    }

    /// Peak level of an output channel since the last call, in dBFS.
    public func takeChannelPeak(deviceUID: String, channel: Int) -> Float {
        guard let engine = outputEngines[deviceUID], channel < engine.channels else { return Level.floorDB }
        let peak = engine.channelPeaks[channel]
        engine.channelPeaks[channel] = 0
        return Level.db(fromGain: peak)
    }

    /// Under- and overruns of a route since it was created (dropouts = audible clicks).
    public func dropouts(_ id: UUID) -> Int {
        guard let renderer = renderers[id] else { return 0 }
        return renderer.totalUnderruns + renderer.totalOverruns
    }

    public struct Diagnostics: Sendable {
        public let underruns: Int
        public let overruns: Int
        public let bufferedFrames: Int
        public let correction: Double
    }

    public func diagnostics(_ id: UUID) -> Diagnostics? {
        guard let renderer = renderers[id] else { return nil }
        return Diagnostics(underruns: renderer.totalUnderruns, overruns: renderer.totalOverruns,
                           bufferedFrames: Int(musubi_ring_available(renderer.feed.ring)), correction: renderer.correction)
    }

    // MARK: - Engines

    private func retire(_ renderer: RouteRenderer) {
        retired.append(renderer)
        Task { [weak self] in
            try? await Task.sleep(for: .seconds(2))
            self?.retired.removeAll { $0 === renderer }
        }
    }

    private func update(_ id: UUID, _ change: (inout Route) -> Void) {
        guard let index = routes.firstIndex(where: { $0.id == id }) else { return }
        change(&routes[index])
    }

    private func applyGates() {
        for route in routes {
            guard case .input = route.source else { continue }
            renderers[route.id]?.gateOpen = talkActive
        }
    }

    private func inputEngine(for uid: String) -> InputEngine? {
        if let engine = inputEngines[uid] { return engine }
        guard let device = devices.first(where: { $0.uid == uid }) else { return nil }
        do {
            let engine = try InputEngine(device: device)
            inputEngines[uid] = engine
            errors[uid] = nil
            return engine
        } catch {
            errors[uid] = "\(device.name) lässt sich nicht öffnen."
            log.error("Input \(device.name, privacy: .public): \(String(describing: error), privacy: .public)")
            return nil
        }
    }

    /// Starts engines for devices with routes, stops the others, and hands each engine its routes.
    private func rebuildOutputs() {
        let wanted = Set(routes.map(\.deviceUID))
        for uid in outputEngines.keys where !wanted.contains(uid) {
            outputEngines.removeValue(forKey: uid)?.stop()
        }
        for uid in wanted {
            if outputEngines[uid] == nil, let device = devices.first(where: { $0.uid == uid }) {
                do {
                    outputEngines[uid] = try OutputEngine(device: device)
                    errors[uid] = nil
                    log.info("Output \(device.name, privacy: .public) started at \(device.sampleRate) Hz")
                } catch {
                    errors[uid] = "\(device.name) lässt sich nicht öffnen."
                    log.error("Output \(device.name, privacy: .public): \(String(describing: error), privacy: .public)")
                }
            }
            outputEngines[uid]?.setRoutes(routes.filter { $0.deviceUID == uid }.compactMap { renderers[$0.id] })
        }
    }

    /// A device was plugged in or out: drop engines of vanished devices, restart those that came back.
    private func devicesChanged() {
        devices = AudioDevices.all()
        let present = Set(devices.map(\.uid))

        for uid in outputEngines.keys where !present.contains(uid) {
            outputEngines.removeValue(forKey: uid)?.stop()
            log.notice("Output device \(uid, privacy: .public) disappeared")
        }
        for uid in inputEngines.keys where !present.contains(uid) {
            inputEngines.removeValue(forKey: uid)?.stop()
            log.notice("Input device \(uid, privacy: .public) disappeared")
        }
        // Reattach input routes of devices that are back (add is idempotent via remove first).
        for route in routes {
            guard case .input(let uid) = route.source, present.contains(uid),
                  let renderer = renderers[route.id], let engine = inputEngine(for: uid) else { continue }
            engine.fanout.remove(renderer.feed)
            engine.fanout.add(renderer.feed)
        }
        rebuildOutputs()
    }
}