downterm/CommandNotch/CommandNotchTests/ScreenRegistryTests.swift

import XCTest
@testable import CommandNotch

@MainActor
final class ScreenRegistryTests: XCTestCase {
    func testRefreshCreatesContextsForConnectedScreensUsingDefaultWorkspace() {
        let workspaceRegistry = makeWorkspaceRegistry()
        let settingsController = makeSettingsController()
        let registry = ScreenRegistry(
            workspaceRegistry: workspaceRegistry,
            settingsController: settingsController,
            connectedScreenIDsProvider: { ["screen-a", "screen-b"] },
            activeScreenIDProvider: { "screen-b" },
            screenLookup: { _ in nil }
        )

        XCTAssertEqual(registry.allScreens().map(\.id), ["screen-a", "screen-b"])
        XCTAssertEqual(
            registry.allScreens().map(\.workspaceID),
            [workspaceRegistry.defaultWorkspaceID, workspaceRegistry.defaultWorkspaceID]
        )
        XCTAssertEqual(registry.activeScreenID(), "screen-b")
    }

    func testAssignWorkspaceUpdatesContextAndSurvivesReconnect() {
        let workspaceRegistry = makeWorkspaceRegistry()
        let settingsController = makeSettingsController()
        var connectedScreenIDs = ["screen-a"]
        let registry = ScreenRegistry(
            workspaceRegistry: workspaceRegistry,
            settingsController: settingsController,
            connectedScreenIDsProvider: { connectedScreenIDs },
            activeScreenIDProvider: { connectedScreenIDs.first },
            screenLookup: { _ in nil }
        )
        let docsWorkspaceID = workspaceRegistry.createWorkspace(named: "Docs")

        registry.assignWorkspace(docsWorkspaceID, to: "screen-a")
        XCTAssertEqual(registry.screenContext(for: "screen-a")?.workspaceID, docsWorkspaceID)

        connectedScreenIDs = []
        registry.refreshConnectedScreens()
        XCTAssertNil(registry.screenContext(for: "screen-a"))

        connectedScreenIDs = ["screen-a"]
        registry.refreshConnectedScreens()
        XCTAssertEqual(registry.screenContext(for: "screen-a")?.workspaceID, docsWorkspaceID)
    }

    func testDeletedWorkspaceAssignmentFallsBackToDefaultWorkspace() {
        let workspaceRegistry = makeWorkspaceRegistry()
        let settingsController = makeSettingsController()
        let assignmentStore = InMemoryScreenAssignmentStore()
        let registry = ScreenRegistry(
            workspaceRegistry: workspaceRegistry,
            settingsController: settingsController,
            assignmentStore: assignmentStore,
            connectedScreenIDsProvider: { ["screen-a"] },
            activeScreenIDProvider: { "screen-a" },
            screenLookup: { _ in nil }
        )
        let reviewWorkspaceID = workspaceRegistry.createWorkspace(named: "Review")

        registry.assignWorkspace(reviewWorkspaceID, to: "screen-a")
        workspaceRegistry.deleteWorkspace(id: reviewWorkspaceID)
        registry.refreshConnectedScreens()

        XCTAssertEqual(
            registry.screenContext(for: "screen-a")?.workspaceID,
            workspaceRegistry.defaultWorkspaceID
        )
        XCTAssertEqual(
            assignmentStore.savedAssignments["screen-a"],
            workspaceRegistry.defaultWorkspaceID
        )
    }

    func testRegistryLoadsPersistedAssignmentsFromStore() {
        let workspaceRegistry = makeWorkspaceRegistry()
        let settingsController = makeSettingsController()
        let docsWorkspaceID = workspaceRegistry.createWorkspace(named: "Docs")
        let assignmentStore = InMemoryScreenAssignmentStore()
        assignmentStore.savedAssignments = ["screen-a": docsWorkspaceID]
        let registry = ScreenRegistry(
            workspaceRegistry: workspaceRegistry,
            settingsController: settingsController,
            assignmentStore: assignmentStore,
            connectedScreenIDsProvider: { ["screen-a"] },
            activeScreenIDProvider: { "screen-a" },
            screenLookup: { _ in nil }
        )

        XCTAssertEqual(registry.screenContext(for: "screen-a")?.workspaceID, docsWorkspaceID)
    }

    func testAssignWorkspacePersistsAssignments() {
        let workspaceRegistry = makeWorkspaceRegistry()
        let settingsController = makeSettingsController()
        let assignmentStore = InMemoryScreenAssignmentStore()
        let registry = ScreenRegistry(
            workspaceRegistry: workspaceRegistry,
            settingsController: settingsController,
            assignmentStore: assignmentStore,
            connectedScreenIDsProvider: { ["screen-a"] },
            activeScreenIDProvider: { "screen-a" },
            screenLookup: { _ in nil }
        )
        let reviewWorkspaceID = workspaceRegistry.createWorkspace(named: "Review")

        registry.assignWorkspace(reviewWorkspaceID, to: "screen-a")

        XCTAssertEqual(assignmentStore.savedAssignments["screen-a"], reviewWorkspaceID)
    }

    func testWorkspaceControllerTracksAssignedWorkspace() {
        let workspaceRegistry = makeWorkspaceRegistry()
        let settingsController = makeSettingsController()
        let registry = ScreenRegistry(
            workspaceRegistry: workspaceRegistry,
            settingsController: settingsController,
            connectedScreenIDsProvider: { ["screen-a"] },
            activeScreenIDProvider: { "screen-a" },
            screenLookup: { _ in nil }
        )
        let docsWorkspaceID = workspaceRegistry.createWorkspace(named: "Docs")

        XCTAssertEqual(registry.workspaceController(for: "screen-a").id, workspaceRegistry.defaultWorkspaceID)

        registry.assignWorkspace(docsWorkspaceID, to: "screen-a")

        XCTAssertEqual(registry.workspaceController(for: "screen-a").id, docsWorkspaceID)
    }

    func testDeleteWorkspaceReassignsConnectedAndPersistedScreensToFallback() {
        let workspaceRegistry = makeWorkspaceRegistry()
        let settingsController = makeSettingsController()
        let docsWorkspaceID = workspaceRegistry.createWorkspace(named: "Docs")
        let reviewWorkspaceID = workspaceRegistry.createWorkspace(named: "Review")
        let assignmentStore = InMemoryScreenAssignmentStore()
        assignmentStore.savedAssignments = ["screen-a": docsWorkspaceID, "screen-b": docsWorkspaceID]
        let registry = ScreenRegistry(
            workspaceRegistry: workspaceRegistry,
            settingsController: settingsController,
            assignmentStore: assignmentStore,
            connectedScreenIDsProvider: { ["screen-a"] },
            activeScreenIDProvider: { "screen-a" },
            screenLookup: { _ in nil }
        )

        registry.assignWorkspace(docsWorkspaceID, to: "screen-a")

        let fallbackWorkspaceID = registry.deleteWorkspace(
            docsWorkspaceID,
            preferredFallback: reviewWorkspaceID
        )

        XCTAssertEqual(fallbackWorkspaceID, reviewWorkspaceID)
        XCTAssertEqual(registry.screenContext(for: "screen-a")?.workspaceID, reviewWorkspaceID)
        XCTAssertEqual(assignmentStore.savedAssignments["screen-a"], reviewWorkspaceID)
        XCTAssertEqual(assignmentStore.savedAssignments["screen-b"], reviewWorkspaceID)
    }

    func testAssignedScreenCountIncludesDisconnectedAssignments() {
        let workspaceRegistry = makeWorkspaceRegistry()
        let settingsController = makeSettingsController()
        let docsWorkspaceID = workspaceRegistry.createWorkspace(named: "Docs")
        let assignmentStore = InMemoryScreenAssignmentStore()
        assignmentStore.savedAssignments = ["screen-a": docsWorkspaceID, "screen-b": docsWorkspaceID]
        let registry = ScreenRegistry(
            workspaceRegistry: workspaceRegistry,
            settingsController: settingsController,
            assignmentStore: assignmentStore,
            connectedScreenIDsProvider: { ["screen-a"] },
            activeScreenIDProvider: { "screen-a" },
            screenLookup: { _ in nil }
        )

        XCTAssertEqual(registry.assignedScreenCount(to: docsWorkspaceID), 2)
        XCTAssertEqual(registry.assignedScreenIDs(to: docsWorkspaceID), ["screen-a", "screen-b"])
    }

    func testClaimWorkspacePresentationTracksPresenterPerWorkspace() {
        let workspaceRegistry = makeWorkspaceRegistry()
        let settingsController = makeSettingsController()
        let docsWorkspaceID = workspaceRegistry.createWorkspace(named: "Docs")
        let registry = ScreenRegistry(
            workspaceRegistry: workspaceRegistry,
            settingsController: settingsController,
            connectedScreenIDsProvider: { ["screen-a", "screen-b"] },
            activeScreenIDProvider: { "screen-a" },
            screenLookup: { _ in nil }
        )

        registry.assignWorkspace(docsWorkspaceID, to: "screen-a")
        registry.assignWorkspace(docsWorkspaceID, to: "screen-b")

        XCTAssertNil(registry.claimWorkspacePresentation(for: "screen-a"))
        XCTAssertEqual(registry.presentingScreenID(for: docsWorkspaceID), "screen-a")
        XCTAssertEqual(registry.claimWorkspacePresentation(for: "screen-b"), "screen-a")
        XCTAssertEqual(registry.presentingScreenID(for: docsWorkspaceID), "screen-b")

        registry.releaseWorkspacePresentation(for: "screen-b")
        XCTAssertNil(registry.presentingScreenID(for: docsWorkspaceID))
    }

    func testAssignWorkspaceReleasesPreviousPresentationOwnership() {
        let workspaceRegistry = makeWorkspaceRegistry()
        let settingsController = makeSettingsController()
        let docsWorkspaceID = workspaceRegistry.createWorkspace(named: "Docs")
        let reviewWorkspaceID = workspaceRegistry.createWorkspace(named: "Review")
        let registry = ScreenRegistry(
            workspaceRegistry: workspaceRegistry,
            settingsController: settingsController,
            connectedScreenIDsProvider: { ["screen-a"] },
            activeScreenIDProvider: { "screen-a" },
            screenLookup: { _ in nil }
        )

        registry.assignWorkspace(docsWorkspaceID, to: "screen-a")
        _ = registry.claimWorkspacePresentation(for: "screen-a")
        XCTAssertEqual(registry.presentingScreenID(for: docsWorkspaceID), "screen-a")

        registry.assignWorkspace(reviewWorkspaceID, to: "screen-a")

        XCTAssertNil(registry.presentingScreenID(for: docsWorkspaceID))
        XCTAssertEqual(registry.screenContext(for: "screen-a")?.workspaceID, reviewWorkspaceID)
    }

    func testConnectedScreenSummariesReflectActiveScreenAndAssignments() {
        let workspaceRegistry = makeWorkspaceRegistry()
        let settingsController = makeSettingsController()
        let docsWorkspaceID = workspaceRegistry.createWorkspace(named: "Docs")
        let registry = ScreenRegistry(
            workspaceRegistry: workspaceRegistry,
            settingsController: settingsController,
            connectedScreenIDsProvider: { ["screen-a", "screen-b"] },
            activeScreenIDProvider: { "screen-b" },
            screenLookup: { _ in nil }
        )

        registry.assignWorkspace(docsWorkspaceID, to: "screen-a")

        XCTAssertEqual(
            registry.connectedScreenSummaries(),
            [
                ConnectedScreenSummary(
                    id: "screen-a",
                    displayName: "Screen 1",
                    isActive: false,
                    assignedWorkspaceID: docsWorkspaceID
                ),
                ConnectedScreenSummary(
                    id: "screen-b",
                    displayName: "Screen 2",
                    isActive: true,
                    assignedWorkspaceID: workspaceRegistry.defaultWorkspaceID
                )
            ]
        )
    }

    private func makeWorkspaceRegistry() -> WorkspaceRegistry {
        let settingsProvider = ScreenRegistryTestSettingsProvider()
        let sessionFactory = ScreenRegistryUnusedTerminalSessionFactory()

        return WorkspaceRegistry(
            initialWorkspaces: [],
            controllerFactory: { summary in
                WorkspaceController(
                    summary: summary,
                    sessionFactory: sessionFactory,
                    settingsProvider: settingsProvider,
                    bootstrapDefaultTab: false
                )
            }
        )
    }

    private func makeSettingsController() -> AppSettingsController {
        AppSettingsController(store: TestAppSettingsStore())
    }
}

private final class InMemoryScreenAssignmentStore: ScreenAssignmentStoreType {
    var savedAssignments: [ScreenID: WorkspaceID] = [:]

    func loadScreenAssignments() -> [ScreenID: WorkspaceID] {
        savedAssignments
    }

    func saveScreenAssignments(_ assignments: [ScreenID: WorkspaceID]) {
        savedAssignments = assignments
    }
}

private final class TestAppSettingsStore: AppSettingsStoreType {
    private var storedSettings = AppSettings.default

    func load() -> AppSettings {
        storedSettings
    }

    func save(_ settings: AppSettings) {
        storedSettings = settings
    }
}

private final class ScreenRegistryTestSettingsProvider: TerminalSessionConfigurationProviding {
    let terminalSessionConfiguration = TerminalSessionConfiguration(fontSize: 13, theme: .terminalApp, shellPath: "")
    let hotkeySettings = AppSettings.default.hotkeys
    let terminalSizePresets = TerminalSizePresetStore.loadDefaults()
}

private struct ScreenRegistryUnusedTerminalSessionFactory: TerminalSessionFactoryType {
    @MainActor
    func makeSession(
        fontSize: CGFloat,
        theme: TerminalTheme,
        shellPath: String,
        initialDirectory: String?
    ) -> TerminalSession {
        fatalError("ScreenRegistryTests should not create live terminal sessions.")
    }
}