import AppKit
import CoreGraphics
import Darwin
import Foundation
import AlignerCore

final class SkyLightSpaceEnumerator: SpaceEnumeratorProtocol {
    enum SkyLightError: Error, CustomStringConvertible {
        case frameworkUnavailable
        case symbolUnavailable(String)
        case managedDisplaySpacesUnavailable

        var description: String {
            switch self {
            case .frameworkUnavailable:
                return "SkyLight framework could not be opened."
            case .symbolUnavailable(let name):
                return "SkyLight symbol unavailable: \(name)."
            case .managedDisplaySpacesUnavailable:
                return "CGSCopyManagedDisplaySpaces returned no usable data."
            }
        }
    }

    private typealias CGSMainConnectionIDFunction = @convention(c) () -> UInt32
    private typealias CGSCopyManagedDisplaySpacesFunction = @convention(c) (UInt32) -> Unmanaged<CFArray>?
    private typealias CGSCopySpacesForWindowsFunction = @convention(c) (UInt32, UInt32, CFArray) -> Unmanaged<CFArray>?

    private let handle: UnsafeMutableRawPointer?
    private let mainConnectionID: CGSMainConnectionIDFunction
    private let copyManagedDisplaySpaces: CGSCopyManagedDisplaySpacesFunction
    private let copySpacesForWindows: CGSCopySpacesForWindowsFunction?

    init() throws {
        guard let handle = dlopen("/System/Library/PrivateFrameworks/SkyLight.framework/SkyLight", RTLD_LAZY) else {
            throw SkyLightError.frameworkUnavailable
        }

        self.handle = handle
        self.mainConnectionID = try Self.symbol("CGSMainConnectionID", in: handle)
        self.copyManagedDisplaySpaces = try Self.symbol("CGSCopyManagedDisplaySpaces", in: handle)
        self.copySpacesForWindows = try? Self.symbol("CGSCopySpacesForWindows", in: handle)
    }

    deinit {
        if let handle {
            dlclose(handle)
        }
    }

    func displays() throws -> [AlignerDisplay] {
        let displayRecords = try managedDisplaySpaceRecords()

        let physicalDisplayUUIDs = Self.physicalDisplayUUIDs()
        let records = displayRecords.flatMap { displayRecord -> [SpaceEnumerationRecord] in
            let displayUUID = stringValue(displayRecord["Display Identifier"])
            let spaceRecords = displayRecord["Spaces"] as? [[String: Any]] ?? []

            return spaceRecords.compactMap { spaceRecord in
                guard let id = uint64Value(spaceRecord["id64"] ?? spaceRecord["id"]) else {
                    return nil
                }

                return SpaceEnumerationRecord(
                    id: id,
                    cgsType: intValue(spaceRecord["type"]) ?? -1,
                    displayUUID: displayUUID
                )
            }
        }

        return SpaceEnumerationPolicy.normalizedDisplays(
            from: records,
            physicalDisplayUUIDs: physicalDisplayUUIDs
        )
    }

    func currentOverlayScreenSpaceStatesByDisplayUUID() throws -> [String: OverlayScreenSpaceState] {
        let displayRecords = try managedDisplaySpaceRecords()
        var result: [String: OverlayScreenSpaceState] = [:]

        for displayRecord in displayRecords {
            let displayUUID = stringValue(displayRecord["Display Identifier"])
            guard !displayUUID.isEmpty,
                  let currentSpace = displayRecord["Current Space"] as? [String: Any]
            else {
                continue
            }

            let isFullscreen = SpaceEnumerationPolicy.spaceType(fromCGSType: intValue(currentSpace["type"]) ?? -1) == .fullscreen
            result[displayUUID] = OverlayScreenSpaceState(
                isFullscreenSpace: isFullscreen,
                isSplitViewSpace: Self.isSplitViewSpace(currentSpace)
            )
        }

        return result
    }

    func currentSpaceIDs() throws -> Set<UInt64> {
        let displayRecords = try managedDisplaySpaceRecords()
        var result = Set<UInt64>()

        for displayRecord in displayRecords {
            guard let currentSpace = displayRecord["Current Space"] as? [String: Any],
                  let id = uint64Value(currentSpace["id64"] ?? currentSpace["id"])
            else {
                continue
            }

            result.insert(id)
        }

        return result
    }

    func spaceIDsByWindowID(windowIDs: [UInt32]) -> [UInt32: [UInt64]] {
        guard let copySpacesForWindows else { return [:] }

        let connectionID = mainConnectionID()
        var result: [UInt32: [UInt64]] = [:]

        for windowID in windowIDs {
            let windowIDArray = [NSNumber(value: windowID)] as CFArray
            guard let unmanagedSpaces = copySpacesForWindows(connectionID, 0x7, windowIDArray) else {
                continue
            }

            let spaces = (unmanagedSpaces.takeRetainedValue() as? [Any] ?? [])
                .compactMap { item -> UInt64? in
                    uint64Value(item)
                }

            if !spaces.isEmpty {
                result[windowID] = spaces
            }
        }

        return result
    }

    private static func symbol<T>(_ name: String, in handle: UnsafeMutableRawPointer) throws -> T {
        guard let rawSymbol = dlsym(handle, name) else {
            throw SkyLightError.symbolUnavailable(name)
        }

        return unsafeBitCast(rawSymbol, to: T.self)
    }

    private func managedDisplaySpaceRecords() throws -> [[String: Any]] {
        let connectionID = mainConnectionID()
        guard
            let unmanagedDisplays = copyManagedDisplaySpaces(connectionID),
            let displayRecords = unmanagedDisplays.takeRetainedValue() as? [[String: Any]]
        else {
            throw SkyLightError.managedDisplaySpacesUnavailable
        }

        return displayRecords
    }

    private static func isSplitViewSpace(_ spaceRecord: [String: Any]) -> Bool {
        guard let tileLayoutManager = spaceRecord["TileLayoutManager"] as? [String: Any] else {
            return false
        }

        if let tileSpaces = tileLayoutManager["TileSpaces"] as? [Any] {
            return tileSpaces.count > 1
        }

        return true
    }

    private static func physicalDisplayUUIDs() -> Set<String> {
        Set(
            NSScreen.screens.compactMap { screen in
                guard
                    let number = screen.deviceDescription[NSDeviceDescriptionKey("NSScreenNumber")] as? NSNumber,
                    let unmanagedUUID = CGDisplayCreateUUIDFromDisplayID(CGDirectDisplayID(number.uint32Value))
                else {
                    return nil
                }

                let uuid = unmanagedUUID.takeRetainedValue()
                return CFUUIDCreateString(nil, uuid) as String?
            }
        )
    }
}

private func stringValue(_ value: Any?) -> String {
    value as? String ?? ""
}

private func intValue(_ value: Any?) -> Int? {
    switch value {
    case let number as NSNumber:
        return number.intValue
    case let int as Int:
        return int
    default:
        return nil
    }
}

private func uint64Value(_ value: Any?) -> UInt64? {
    switch value {
    case let number as NSNumber:
        return number.uint64Value
    case let uint64 as UInt64:
        return uint64
    case let uint as UInt:
        return UInt64(uint)
    case let int as Int where int >= 0:
        return UInt64(int)
    default:
        return nil
    }
}