Musubi/App/Sources/OutputController.swift

import AppKit
import MusubiCore
import NDIKit
import Observation
import VideoOutput

/// A connected screen as shown in the UI.
struct ScreenInfo: Identifiable, Equatable {
    let id: CGDirectDisplayID
    let name: String
    let width: Int
    let height: Int
    var refreshRate: Double
}

/// One source shown fullscreen on one screen.
@MainActor
final class ActiveOutput {
    let source: Source
    let receiver: NDIReceiver
    let window: OutputWindow
    let displayID: CGDirectDisplayID
    let originalMode: CGDisplayMode?
    var snapshot: NDIReceiver.Snapshot

    init(source: Source, receiver: NDIReceiver, window: OutputWindow, displayID: CGDirectDisplayID, originalMode: CGDisplayMode?) {
        self.source = source
        self.receiver = receiver
        self.window = window
        self.displayID = displayID
        self.originalMode = originalMode
        self.snapshot = receiver.snapshot
    }
}

@MainActor @Observable
final class OutputController {
    static let targetRefreshRate = 50.0

    private(set) var screens: [ScreenInfo] = []
    private(set) var outputs: [CGDirectDisplayID: ActiveOutput] = [:]
    /// Bumped once per second so views re-read the receiver snapshots.
    private(set) var tick = 0

    private var sleepBlocker: SleepBlocker?
    private var timer: Timer?
    private let recorder = StatsRecorder()

    var isSleepBlocked: Bool { sleepBlocker != nil }

    init() {
        refreshScreens()
        NotificationCenter.default.addObserver(forName: NSApplication.didChangeScreenParametersNotification, object: nil, queue: .main) { [weak self] _ in
            MainActor.assumeIsolated { self?.screensChanged() }
        }
        timer = Timer.scheduledTimer(withTimeInterval: 1, repeats: true) { [weak self] _ in
            MainActor.assumeIsolated { self?.updateSnapshots() }
        }
    }

    func show(_ source: Source, on displayID: CGDirectDisplayID) {
        guard let screen = NSScreen.screens.first(where: { $0.displayID == displayID }) else { return }
        stop(displayID)

        let window = OutputWindow(screen: screen)
        let sink = window.sink
        let receiverName = "Musubi \(screen.localizedName)"
        guard let receiver = NDIReceiver(source: source, receiverName: receiverName, onFrame: { sink.display($0) }) else {
            Log.ndi.error("Could not create receiver for \(source.ndiName, privacy: .public)")
            return
        }

        let mode = DisplayModeSwitcher.setRefreshRate(Self.targetRefreshRate, on: displayID)
        window.onCancel = { [weak self] in self?.stop(displayID) }
        window.setFrame(screen.frame, display: true)
        window.orderFrontRegardless()

        outputs[displayID] = ActiveOutput(source: source, receiver: receiver, window: window, displayID: displayID, originalMode: mode.original)
        Log.ndi.info("Output \(source.ndiName, privacy: .public) → \(screen.localizedName, privacy: .public) @ \(mode.refreshRate) Hz")
        updateSleepBlocker()
        refreshScreens()
    }

    func stop(_ displayID: CGDirectDisplayID) {
        guard let output = outputs.removeValue(forKey: displayID) else { return }
        output.receiver.stop()
        output.window.orderOut(nil)
        if let original = output.originalMode {
            DisplayModeSwitcher.restore(original, on: displayID)
        }
        Log.ndi.info("Output stopped on display \(displayID)")
        updateSleepBlocker()
        refreshScreens()
    }

    /// Applies `--route "<NDI name>=<screen index>"` launch arguments.
    func applyLaunchRoutes(_ arguments: [String] = ProcessInfo.processInfo.arguments) {
        for route in RouteArgument.parse(arguments) {
            guard route.screenIndex < NSScreen.screens.count else {
                Log.ndi.error("--route: no screen \(route.screenIndex)")
                continue
            }
            show(Source(ndiName: route.ndiName), on: NSScreen.screens[route.screenIndex].displayID)
        }
    }

    private func updateSnapshots() {
        for output in outputs.values {
            output.snapshot = output.receiver.snapshot
        }
        recorder.recordIfDue(outputs: Array(outputs.values), screens: screens)
        tick += 1
    }

    private func updateSleepBlocker() {
        if outputs.isEmpty {
            sleepBlocker = nil
        } else if sleepBlocker == nil {
            sleepBlocker = SleepBlocker(reason: "Musubi: Ausgabe aktiv")
        }
    }

    private func screensChanged() {
        let connected = Set(NSScreen.screens.map(\.displayID))
        for displayID in outputs.keys where !connected.contains(displayID) {
            Log.ndi.notice("Display \(displayID) disconnected, stopping its output")
            stop(displayID)
        }
        // Keep windows covering their screen after resolution changes.
        for output in outputs.values {
            if let screen = NSScreen.screens.first(where: { $0.displayID == output.displayID }) {
                output.window.setFrame(screen.frame, display: true)
            }
        }
        refreshScreens()
    }

    private func refreshScreens() {
        screens = NSScreen.screens.map { screen in
            let displayID = screen.displayID
            return ScreenInfo(
                id: displayID,
                name: screen.localizedName,
                width: Int(CGDisplayPixelsWide(displayID)),
                height: Int(CGDisplayPixelsHigh(displayID)),
                refreshRate: DisplayModeSwitcher.currentRefreshRate(of: displayID)
            )
        }
    }
}

extension NSScreen {
    var displayID: CGDirectDisplayID {
        (deviceDescription[NSDeviceDescriptionKey("NSScreenNumber")] as? NSNumber)?.uint32Value ?? 0
    }
}