Musubi/App/Sources/OutputController.swift

import AppKit
import MusubiCore
import NDIKit
import Remote
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
    /// Mode before Musubi switched to 50 Hz; restored when the output ends.
    var 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] = [:]
    /// What should be shown. Assignments for screens that are not connected stay here and start
    /// as soon as the screen appears (monitors often come up late after a reboot).
    private(set) var desired: [VideoAssignment] = []
    /// Called after the desired configuration changed (for autosave and remote status).
    var onChange: (() -> Void)?
    /// 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
            // Deferred: the notification can arrive synchronously inside a display mode switch.
            DispatchQueue.main.async { MainActor.assumeIsolated { self?.screensChanged() } }
        }
        timer = Timer.scheduledTimer(withTimeInterval: 1, repeats: true) { [weak self] _ in
            MainActor.assumeIsolated { self?.updateSnapshots() }
        }
    }

    /// Shows `source` on a screen and remembers it.
    func show(_ source: Source, on displayID: CGDirectDisplayID) {
        guard let ref = screenRef(for: displayID) else { return }
        show(VideoAssignment(screen: ref, source: source))
    }

    func show(_ assignment: VideoAssignment) {
        desired.removeAll { $0.screen == assignment.screen }
        desired.append(assignment)
        reconcile()
        onChange?()
    }

    /// Stops the output on a screen and forgets it.
    func stop(_ displayID: CGDirectDisplayID) {
        guard let ref = screenRef(for: displayID) else { return halt(displayID) }
        stop(ref)
    }

    func stop(_ ref: ScreenRef) {
        let names = NSScreen.screens.map(\.localizedName)
        let index = ScreenMatching.index(of: ref, in: names)
        desired.removeAll { $0.screen == ref || (index != nil && ScreenMatching.index(of: $0.screen, in: names) == index) }
        reconcile()
        onChange?()
    }

    /// Replaces the whole video configuration (scene load, remote apply, restore).
    func apply(_ video: [VideoAssignment]) {
        desired = video
        reconcile()
        onChange?()
    }

    private var reconciling = false
    private var reconcileAgain = false

    /// Starts and stops outputs so that connected screens match `desired`.
    /// Never nested: display mode switches post screen-change notifications while we are in here.
    private func reconcile() {
        guard !reconciling else {
            reconcileAgain = true
            return
        }
        reconciling = true
        defer { reconciling = false }
        repeat {
            reconcileAgain = false
            reconcileOnce()
        } while reconcileAgain
    }

    private func reconcileOnce() {
        let screens = NSScreen.screens
        let resolved = ScreenMatching.resolve(desired, in: screens.map(\.localizedName))
        let wanted = Dictionary(uniqueKeysWithValues: resolved.map { (screens[$0.key].displayID, $0.value) })
        let current = outputs.mapValues(\.source)
        let plan = Reconcile.plan(current: current, desired: wanted)
        plan.stop.forEach(halt)
        for (displayID, source) in plan.start {
            start(source, on: displayID)
        }
    }

    func screenRef(for displayID: CGDirectDisplayID) -> ScreenRef? {
        guard let index = NSScreen.screens.firstIndex(where: { $0.displayID == displayID }) else { return nil }
        return ScreenRef(name: NSScreen.screens[index].localizedName, index: index)
    }

    /// Screens with their outputs, for the remote status.
    func screenStatuses() -> [ScreenStatus] {
        screens.compactMap { screen in
            guard let ref = screenRef(for: screen.id) else { return nil }
            let output = outputs[screen.id].map { output -> OutputStatus in
                let s = output.snapshot
                let state: OutputStatus.State = switch s.state {
                case .connecting: .connecting
                case .live: .live
                case .stalled: .stalled
                }
                return OutputStatus(source: output.source, state: state, width: s.width, height: s.height, frameRate: s.frameRate,
                                    lostFrames: s.stats.missing + s.stats.droppedByReceiver, lossRate: s.stats.lossRate)
            }
            return ScreenStatus(screen: ref, width: screen.width, height: screen.height, refreshRate: screen.refreshRate, output: output)
        }
    }

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

        // Switching sources on a screen keeps its 50 Hz mode instead of bouncing to 60 Hz and back.
        let previous = outputs.removeValue(forKey: displayID)
        previous?.receiver.stop()

        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)")
            previous?.window.orderOut(nil)
            if let original = previous?.originalMode { DisplayModeSwitcher.restore(original, on: displayID) }
            return
        }
        window.onCancel = { [weak self] in self?.stop(displayID) }
        window.setFrame(screen.frame, display: true)
        window.orderFrontRegardless()
        previous?.window.orderOut(nil)

        // Register before the mode switch: it posts a screen-change notification.
        let output = ActiveOutput(source: source, receiver: receiver, window: window, displayID: displayID, originalMode: previous?.originalMode)
        outputs[displayID] = output
        let mode = DisplayModeSwitcher.setRefreshRate(Self.targetRefreshRate, on: displayID)
        if output.originalMode == nil { output.originalMode = mode.original }
        window.setFrame(screen.frame, display: true)

        Log.ndi.info("Output \(source.ndiName, privacy: .public) → \(screen.localizedName, privacy: .public) @ \(mode.refreshRate) Hz")
        updateSleepBlocker()
        refreshScreens()
    }

    private func halt(_ 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() {
        // 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)
            }
        }
        let connected = Set(NSScreen.screens.map(\.displayID))
        for displayID in outputs.keys where !connected.contains(displayID) {
            Log.ndi.notice("Display \(displayID) disconnected, output paused until it returns")
            halt(displayID)
        }
        // Screens that came (back) get their remembered output.
        reconcile()
        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
    }
}