WorkspacePool.swift

import Dependencies
import Foundation
import Logger
import XcodeInspector

public struct WorkspacePoolDependencyKey: DependencyKey {
    public static var liveValue: WorkspacePool = .init()
}

public extension DependencyValues {
    var workspacePool: WorkspacePool {
        get { self[WorkspacePoolDependencyKey.self] }
        set { self[WorkspacePoolDependencyKey.self] = newValue }
    }
}

@globalActor public enum WorkspaceActor {
    public actor TheActor {}
    public static let shared = TheActor()
}

public class WorkspacePool {
    public enum Error: Swift.Error, LocalizedError {
        case invalidWorkspaceURL(URL)

        public var errorDescription: String? {
            switch self {
            case let .invalidWorkspaceURL(url):
                return "Invalid workspace URL: \(url)"
            }
        }
    }

    public internal(set) var workspaces: [URL: Workspace] = [:]
    var plugins = [ObjectIdentifier: (Workspace) -> WorkspacePlugin]()

    public init(
        workspaces: [URL: Workspace] = [:],
        plugins: [ObjectIdentifier: (Workspace) -> WorkspacePlugin] = [:]
    ) {
        self.workspaces = workspaces
        self.plugins = plugins
    }

    public func registerPlugin<Plugin: WorkspacePlugin>(_ plugin: @escaping (Workspace) -> Plugin) {
        let id = ObjectIdentifier(Plugin.self)
        let erasedPlugin: (Workspace) -> WorkspacePlugin = { plugin($0) }
        plugins[id] = erasedPlugin

        for workspace in workspaces.values {
            addPlugin(erasedPlugin, id: id, to: workspace)
        }
    }

    public func unregisterPlugin<Plugin: WorkspacePlugin>(_: Plugin.Type) {
        let id = ObjectIdentifier(Plugin.self)
        plugins[id] = nil

        for workspace in workspaces.values {
            removePlugin(id: id, from: workspace)
        }
    }

    public func fetchFilespaceIfExisted(fileURL: URL) -> Filespace? {
        let filespaces = workspaces.values.compactMap { $0.filespaces[fileURL] }
        if filespaces.isEmpty { return nil }
        if filespaces.count == 1 { return filespaces.first }
        Logger.workspacePool.info("Multiple workspaces found with file: \(fileURL)")
        // If multiple workspaces are found, return the first with a suggestion
        return filespaces.first { $0.presentingSuggestion != nil } ?? filespaces.first { $0.presentingNESSuggestion != nil }
    }
    
    public func fetchWorkspaceAndFilespace(fileURL: URL) -> (Workspace, Filespace)? {
        var workspace: Workspace?
        var filespace: Filespace?
        
        for wp in workspaces.values {
            if let fp = wp.filespaces[fileURL] {
                if fp.presentingSuggestion != nil || fp.presentingNESSuggestion != nil {
                    return (wp, fp)
                }
                workspace = wp
                filespace = fp
            }
        }
        
        return workspace.flatMap { ws in filespace.map { fs in (ws, fs) } }
    }

    @WorkspaceActor
    public func fetchOrCreateWorkspace(workspaceURL: URL) async throws -> Workspace {
        guard workspaceURL != URL(fileURLWithPath: "/") else {
            throw Error.invalidWorkspaceURL(workspaceURL)
        }

        if let existed = workspaces[workspaceURL] {
            return existed
        }

        let new = createNewWorkspace(workspaceURL: workspaceURL)
        workspaces[workspaceURL] = new
        return new
    }

    @WorkspaceActor
    public func fetchOrCreateWorkspaceAndFilespace(fileURL: URL) async throws
        -> (workspace: Workspace, filespace: Filespace)
    {
        // If we can get the workspace URL directly.
        if let currentWorkspaceURL = await XcodeInspector.shared.safe.realtimeActiveWorkspaceURL {
            if let existed = workspaces[currentWorkspaceURL] {
                // Reuse the existed workspace.
                let filespace = try await existed.createFilespaceIfNeeded(fileURL: fileURL)
                return (existed, filespace)
            }

            let new = createNewWorkspace(workspaceURL: currentWorkspaceURL)
            workspaces[currentWorkspaceURL] = new
            let filespace = try await new.createFilespaceIfNeeded(fileURL: fileURL)
            return (new, filespace)
        }

        // If not, we try to reuse a filespace if found.
        //
        // Sometimes, we can't get the project root path from Xcode window, for example, when the
        // quick open window in displayed.
        for workspace in workspaces.values {
            if let filespace = workspace.filespaces[fileURL] {
                return (workspace, filespace)
            }
        }

        // If we can't find the workspace URL, we will try to guess it.
        // Most of the time we won't enter this branch, just incase.

        if let workspaceURL = WorkspaceXcodeWindowInspector.extractProjectURL(
            workspaceURL: nil,
            documentURL: fileURL
        ) {
            let workspace = {
                if let existed = workspaces[workspaceURL] {
                    return existed
                }
                // Reuse existed workspace if possible
                for (_, workspace) in workspaces {
                    if fileURL.path.hasPrefix(workspace.projectRootURL.path) {
                        return workspace
                    }
                }
                return createNewWorkspace(workspaceURL: workspaceURL)
            }()

            let filespace = try await workspace.createFilespaceIfNeeded(fileURL: fileURL)
            workspaces[workspaceURL] = workspace
            workspace.refreshUpdateTime()
            return (workspace, filespace)
        }
        
        throw Workspace.CantFindWorkspaceError()
    }

    @WorkspaceActor
    public func removeWorkspace(url: URL) {
        workspaces[url] = nil
    }
}

extension WorkspacePool {
    func addPlugin(
        _ plugin: (Workspace) -> WorkspacePlugin,
        id: ObjectIdentifier,
        to workspace: Workspace
    ) {
        if workspace.plugins[id] != nil { return }
        workspace.plugins[id] = plugin(workspace)
    }

    func removePlugin(id: ObjectIdentifier, from workspace: Workspace) {
        workspace.plugins[id] = nil
    }

    func createNewWorkspace(workspaceURL: URL) -> Workspace {
        let new = Workspace(workspaceURL: workspaceURL)
        for (id, plugin) in plugins {
            addPlugin(plugin, id: id, to: new)
        }
        return new
    }
}