downterm/Downterm/CommandNotch/Models/ScreenRegistry.swift

import AppKit
import Combine
import SwiftUI

struct ConnectedScreenSummary: Identifiable, Equatable {
    let id: ScreenID
    let displayName: String
    let isActive: Bool
    let assignedWorkspaceID: WorkspaceID
}

@MainActor
final class ScreenRegistry: ObservableObject {
    static let shared = ScreenRegistry(assignmentStore: UserDefaultsScreenAssignmentStore())

    @Published private(set) var screenContexts: [ScreenContext] = []

    private let workspaceRegistry: WorkspaceRegistry
    private let settingsController: AppSettingsController
    private let assignmentStore: any ScreenAssignmentStoreType
    private let connectedScreenIDsProvider: @MainActor () -> [ScreenID]
    private let activeScreenIDProvider: @MainActor () -> ScreenID?
    private let screenLookup: @MainActor (ScreenID) -> NSScreen?

    private var contextsByID: [ScreenID: ScreenContext] = [:]
    private var preferredAssignments: [ScreenID: WorkspaceID]
    private var workspacePresenters: [WorkspaceID: ScreenID] = [:]
    private var cancellables = Set<AnyCancellable>()

    init(
        workspaceRegistry: WorkspaceRegistry? = nil,
        settingsController: AppSettingsController? = nil,
        assignmentStore: (any ScreenAssignmentStoreType)? = nil,
        initialAssignments: [ScreenID: WorkspaceID]? = nil,
        connectedScreenIDsProvider: @escaping @MainActor () -> [ScreenID] = {
            NSScreen.screens.map(\.displayUUID)
        },
        activeScreenIDProvider: @escaping @MainActor () -> ScreenID? = {
            let mouseLocation = NSEvent.mouseLocation
            return NSScreen.screens.first { NSMouseInRect(mouseLocation, $0.frame, false) }?.displayUUID
                ?? NSScreen.main?.displayUUID
        },
        screenLookup: @escaping @MainActor (ScreenID) -> NSScreen? = { screenID in
            NSScreen.screens.first { $0.displayUUID == screenID }
        }
    ) {
        let resolvedWorkspaceRegistry = workspaceRegistry ?? WorkspaceRegistry.shared
        let resolvedSettingsController = settingsController ?? AppSettingsController.shared
        let resolvedAssignmentStore = assignmentStore ?? UserDefaultsScreenAssignmentStore()

        self.workspaceRegistry = resolvedWorkspaceRegistry
        self.settingsController = resolvedSettingsController
        self.assignmentStore = resolvedAssignmentStore
        self.preferredAssignments = initialAssignments ?? resolvedAssignmentStore.loadScreenAssignments()
        self.connectedScreenIDsProvider = connectedScreenIDsProvider
        self.activeScreenIDProvider = activeScreenIDProvider
        self.screenLookup = screenLookup

        observeWorkspaceChanges()
        refreshConnectedScreens()
    }

    func allScreens() -> [ScreenContext] {
        screenContexts
    }

    func screenContext(for id: ScreenID) -> ScreenContext? {
        contextsByID[id]
    }

    func workspaceController(for screenID: ScreenID) -> WorkspaceController {
        let workspaceID = contextsByID[screenID]?.workspaceID ?? workspaceRegistry.defaultWorkspaceID
        return workspaceRegistry.controller(for: workspaceID) ?? workspaceRegistry.defaultWorkspaceController
    }

    func assignedScreenIDs(to workspaceID: WorkspaceID) -> [ScreenID] {
        preferredAssignments
            .filter { $0.value == workspaceID }
            .map(\.key)
            .sorted()
    }

    func assignedScreenCount(to workspaceID: WorkspaceID) -> Int {
        assignedScreenIDs(to: workspaceID).count
    }

    func connectedScreenSummaries() -> [ConnectedScreenSummary] {
        let activeScreenID = activeScreenID()

        return screenContexts.enumerated().map { index, context in
            ConnectedScreenSummary(
                id: context.id,
                displayName: resolvedDisplayName(for: context.id, fallbackIndex: index),
                isActive: context.id == activeScreenID,
                assignedWorkspaceID: context.workspaceID
            )
        }
    }

    func assignWorkspace(_ workspaceID: WorkspaceID, to screenID: ScreenID) {
        guard workspaceRegistry.controller(for: workspaceID) != nil else { return }

        let previousWorkspaceID = contextsByID[screenID]?.workspaceID ?? preferredAssignments[screenID]
        preferredAssignments[screenID] = workspaceID
        contextsByID[screenID]?.updateWorkspace(id: workspaceID)

        if let previousWorkspaceID,
           previousWorkspaceID != workspaceID,
           workspacePresenters[previousWorkspaceID] == screenID {
            workspacePresenters.removeValue(forKey: previousWorkspaceID)
        }

        persistAssignments()
    }

    @discardableResult
    func assignActiveScreen(to workspaceID: WorkspaceID) -> ScreenID? {
        guard let screenID = activeScreenID() else { return nil }
        assignWorkspace(workspaceID, to: screenID)
        return screenID
    }

    func presentingScreenID(for workspaceID: WorkspaceID) -> ScreenID? {
        guard let screenID = workspacePresenters[workspaceID] else { return nil }
        guard preferredAssignments[screenID] == workspaceID else {
            workspacePresenters.removeValue(forKey: workspaceID)
            return nil
        }
        return screenID
    }

    @discardableResult
    func claimWorkspacePresentation(for screenID: ScreenID) -> ScreenID? {
        guard let workspaceID = contextsByID[screenID]?.workspaceID ?? preferredAssignments[screenID] else {
            return nil
        }

        let previousPresenter = workspacePresenters[workspaceID]
        workspacePresenters[workspaceID] = screenID
        return previousPresenter == screenID ? nil : previousPresenter
    }

    func releaseWorkspacePresentation(for screenID: ScreenID) {
        workspacePresenters = workspacePresenters.filter { $0.value != screenID }
    }

    @discardableResult
    func deleteWorkspace(
        _ workspaceID: WorkspaceID,
        preferredFallback preferredFallbackID: WorkspaceID? = nil
    ) -> WorkspaceID? {
        guard workspaceRegistry.canDeleteWorkspace(id: workspaceID) else {
            return nil
        }

        guard let fallbackWorkspaceID = workspaceRegistry.deletionFallbackWorkspaceID(
            forDeleting: workspaceID,
            preferredFallback: preferredFallbackID
        ) else {
            return nil
        }

        workspacePresenters.removeValue(forKey: workspaceID)

        for (screenID, assignedWorkspaceID) in preferredAssignments where assignedWorkspaceID == workspaceID {
            preferredAssignments[screenID] = fallbackWorkspaceID
        }

        for context in contextsByID.values where context.workspaceID == workspaceID {
            context.updateWorkspace(id: fallbackWorkspaceID)
        }

        guard workspaceRegistry.deleteWorkspace(id: workspaceID) else {
            return nil
        }

        persistAssignments()
        return fallbackWorkspaceID
    }

    func activeScreenID() -> ScreenID? {
        activeScreenIDProvider() ?? screenContexts.first?.id
    }

    func refreshConnectedScreens() {
        let connectedScreenIDs = connectedScreenIDsProvider()
        let validWorkspaceIDs = Set(workspaceRegistry.allWorkspaceSummaries().map(\.id))
        let defaultWorkspaceID = workspaceRegistry.defaultWorkspaceID
        var nextContextsByID: [ScreenID: ScreenContext] = [:]
        var nextContexts: [ScreenContext] = []

        for screenID in connectedScreenIDs {
            let workspaceID = resolvedWorkspaceID(
                for: screenID,
                validWorkspaceIDs: validWorkspaceIDs,
                defaultWorkspaceID: defaultWorkspaceID
            )

            let context = contextsByID[screenID] ?? ScreenContext(
                id: screenID,
                workspaceID: workspaceID,
                settingsController: settingsController,
                screenProvider: screenLookup
            )

            context.updateWorkspace(id: workspaceID)
            context.refreshClosedSize()

            nextContextsByID[screenID] = context
            nextContexts.append(context)
        }

        contextsByID = nextContextsByID
        screenContexts = nextContexts
        reconcileWorkspacePresenters()
        persistAssignments()
    }

    private func resolvedWorkspaceID(
        for screenID: ScreenID,
        validWorkspaceIDs: Set<WorkspaceID>,
        defaultWorkspaceID: WorkspaceID
    ) -> WorkspaceID {
        guard let preferredWorkspaceID = preferredAssignments[screenID],
              validWorkspaceIDs.contains(preferredWorkspaceID) else {
            preferredAssignments[screenID] = defaultWorkspaceID
            return defaultWorkspaceID
        }

        return preferredWorkspaceID
    }

    private func observeWorkspaceChanges() {
        workspaceRegistry.$workspaceSummaries
            .dropFirst()
            .sink { [weak self] _ in
                Task { @MainActor [weak self] in
                    self?.refreshConnectedScreens()
                }
            }
            .store(in: &cancellables)
    }

    private func persistAssignments() {
        assignmentStore.saveScreenAssignments(preferredAssignments)
    }

    private func reconcileWorkspacePresenters() {
        let validScreenIDs = Set(contextsByID.keys)
        let validAssignments = preferredAssignments

        workspacePresenters = workspacePresenters.filter { workspaceID, screenID in
            validScreenIDs.contains(screenID) && validAssignments[screenID] == workspaceID
        }
    }

    private func resolvedDisplayName(for screenID: ScreenID, fallbackIndex: Int) -> String {
        let fallbackName = "Screen \(fallbackIndex + 1)"
        guard let screen = screenLookup(screenID) else {
            return fallbackName
        }

        let localizedName = screen.localizedName.trimmingCharacters(in: .whitespacesAndNewlines)
        return localizedName.isEmpty ? fallbackName : localizedName
    }
}

extension ScreenRegistry: ScreenRegistryType {}