smart_printer/smart_printer/ScannerService.swift

import Cocoa
import ImageCaptureCore

// MARK: - Scanner Service

@MainActor
final class ScannerService: NSObject {
    static let shared = ScannerService()

    enum State: Equatable {
        case idle
        case browsing
        case connecting
        case ready
        case scanning
        case error(String)
    }

    private(set) var state: State = .idle {
        didSet { onStateChanged?(state) }
    }

    private(set) var localDevices: [ICScannerDevice] = [] {
        didSet { onDevicesChanged?() }
    }

    private(set) var sharedDevices: [ICScannerDevice] = [] {
        didSet { onDevicesChanged?() }
    }

    private(set) var selectedDevice: ICScannerDevice? {
        didSet { onSelectionChanged?() }
    }

    var paperSize: PaperSize = AppSettings.defaultPaperSize
    var resolution: ScanResolution = AppSettings.scanResolution

    var onStateChanged: ((State) -> Void)?
    var onDevicesChanged: (() -> Void)?
    var onSelectionChanged: (() -> Void)?
    var onScanComplete: ((Result<NSImage, Error>) -> Void)?

    private let browser = ICDeviceBrowser()
    private var scanCompletion: ((Result<NSImage, Error>) -> Void)?
    private var shouldScanWhenReady = false
    private var scannedImageURL: URL?
    private var didFinishCurrentScan = false
    private var accumulatedBands: [ICScannerBandData] = []
    private var isClosingSession = false
    private var isOpeningSession = false
    private var isMaintainingSession = false
    private var hasEstablishedSession = false
    private var lastSelectedDeviceName: String?
    private var pendingPreferredDeviceName: String?
    private var sessionCloseContinuation: CheckedContinuation<Void, Never>?
    private let connectionTimeoutSeconds: UInt64 = 20
    private let sessionCloseTimeoutSeconds: UInt64 = 6

    private override init() {
        super.init()
        browser.delegate = self
        let maskRaw = ICDeviceTypeMask.scanner.rawValue
            | ICDeviceLocationTypeMask.local.rawValue
            | ICDeviceLocationTypeMask.shared.rawValue
            | ICDeviceLocationTypeMask.bonjour.rawValue
        browser.browsedDeviceTypeMask = ICDeviceTypeMask(rawValue: maskRaw) ?? .scanner
    }

    /// Opens the scanner panel. Reuses an existing session when one is already open.
    func prepareForScanning() async {
        clearConnectionTimeout()
        isOpeningSession = false
        isClosingSession = false
        didFinishCurrentScan = false
        shouldScanWhenReady = false
        scanCompletion = nil
        scannedImageURL = nil
        accumulatedBands = []

        if let device = selectedDevice {
            device.delegate = self
            if device.hasOpenSession {
                state = .ready
                ensureBrowsing()
                return
            }
            if hasEstablishedSession {
                state = .ready
                ensureBrowsing()
                maintainSession(for: device)
                return
            }
        }

        if let device = deviceMatchingLastSelection() {
            selectedDevice = device
            device.delegate = self
            if device.hasOpenSession {
                state = .ready
                ensureBrowsing()
                return
            }
            if hasEstablishedSession {
                state = .ready
                ensureBrowsing()
                maintainSession(for: device)
                return
            }
        }

        pendingPreferredDeviceName = lastSelectedDeviceName
        ensureBrowsing()
        state = .browsing
    }

    /// Panel closed — keep the scanner session alive for the next open.
    func detachFromUI() {
        clearConnectionTimeout()
        scanCompletion = nil
        shouldScanWhenReady = false
    }

    /// Fully releases the scanner when leaving the scan feature entirely.
    func releaseScanner() async {
        clearConnectionTimeout()
        detachFromUI()

        if let device = selectedDevice {
            await closeSession(on: device)
        }

        if browser.isBrowsing {
            browser.stop()
        }

        resetAfterFullShutdown()
    }

    /// True when reconnecting after a previous successful scan in this session.
    var isWarmReconnecting: Bool {
        isMaintainingSession || (isOpeningSession && hasEstablishedSession)
    }

    func selectDevice(_ device: ICScannerDevice) {
        let name = device.name ?? device.locationDescription
        if let name { lastSelectedDeviceName = name }
        pendingPreferredDeviceName = nil

        if selectedDevice === device {
            if device.hasOpenSession {
                device.delegate = self
                state = .ready
            } else if case .connecting = state {
                return
            } else {
                Task { await openSession(on: device, warm: hasEstablishedSession) }
            }
            return
        }

        selectedDevice = device

        if device.hasOpenSession {
            device.delegate = self
            state = .ready
        } else {
            Task { await openSession(on: device, warm: hasEstablishedSession) }
        }
    }

    func scan() {
        guard let device = selectedDevice else { return }
        scanCompletion = onScanComplete
        scannedImageURL = nil
        didFinishCurrentScan = false
        accumulatedBands = []

        if case .connecting = state {
            shouldScanWhenReady = true
            return
        }

        device.delegate = self
        if device.hasOpenSession {
            configureFunctionalUnit(on: device)
            state = .scanning
            device.requestScan()
        } else {
            shouldScanWhenReady = true
            Task { await openSession(on: device, warm: hasEstablishedSession) }
        }
    }

    func cancelScan() {
        selectedDevice?.cancelScan()
        scanCompletion?(.failure(ScanError.cancelled))
        scanCompletion = nil
        scannedImageURL = nil
        didFinishCurrentScan = true
        accumulatedBands = []
        state = selectedDevice?.hasOpenSession == true ? .ready : .browsing
    }

    private func ensureBrowsing() {
        if !browser.isBrowsing {
            browser.start()
        }
    }

    private func deviceMatchingLastSelection() -> ICScannerDevice? {
        guard let name = lastSelectedDeviceName else { return nil }
        return (localDevices + sharedDevices).first {
            ($0.name ?? $0.locationDescription) == name
        }
    }

    private func resetAfterFullShutdown() {
        isClosingSession = false
        isOpeningSession = false
        isMaintainingSession = false
        hasEstablishedSession = false
        pendingPreferredDeviceName = nil
        selectedDevice = nil
        localDevices = []
        sharedDevices = []
        state = .idle
        scannedImageURL = nil
        didFinishCurrentScan = false
        accumulatedBands = []
        shouldScanWhenReady = false
    }

    private func closeSession(on device: ICScannerDevice) async {
        device.delegate = nil
        if state == .scanning {
            device.cancelScan()
        }
        guard device.hasOpenSession else {
            isClosingSession = false
            return
        }
        isClosingSession = true
        do {
            try await device.requestCloseSession()
            processSessionClosed(on: device, error: nil)
        } catch {
            processSessionClosed(on: device, error: error)
        }
        await waitForSessionToCloseAsync()
    }

    private func waitForSessionToCloseAsync() async {
        if !isClosingSession, selectedDevice?.hasOpenSession != true {
            return
        }
        await withCheckedContinuation { (continuation: CheckedContinuation<Void, Never>) in
            sessionCloseContinuation = continuation
            Task { @MainActor in
                try? await Task.sleep(nanoseconds: sessionCloseTimeoutSeconds * 1_000_000_000)
                resumeSessionCloseWait()
            }
        }
    }

    private func resumeSessionCloseWait() {
        guard let continuation = sessionCloseContinuation else { return }
        sessionCloseContinuation = nil
        isClosingSession = false
        continuation.resume()
    }

    private func maintainSession(for device: ICScannerDevice) {
        guard selectedDevice === device, !device.hasOpenSession else { return }
        guard !isOpeningSession, !isMaintainingSession, !isClosingSession else { return }
        isMaintainingSession = true
        Task {
            await openSession(on: device, warm: true)
            isMaintainingSession = false
        }
    }

    private func openSession(on device: ICScannerDevice, warm: Bool = false) async {
        guard selectedDevice === device else { return }

        clearConnectionTimeout()
        device.delegate = self

        if device.hasOpenSession {
            isOpeningSession = false
            state = .ready
            if shouldScanWhenReady {
                shouldScanWhenReady = false
                configureFunctionalUnit(on: device)
                state = .scanning
                device.requestScan()
            }
            return
        }

        isOpeningSession = true
        state = (warm && hasEstablishedSession) ? .ready : .connecting
        scheduleConnectionTimeout(for: device)

        do {
            try await device.requestOpenSession()
            processSessionOpened(on: device, error: nil)
        } catch {
            processSessionOpened(on: device, error: error)
        }
    }

    private func processSessionOpened(on scanner: ICScannerDevice, error: Error?) {
        guard selectedDevice === scanner else { return }
        guard isOpeningSession || state == .connecting else { return }

        clearConnectionTimeout()
        isOpeningSession = false

        if let error {
            state = .error(error.localizedDescription)
            finishScan(with: .failure(error))
            return
        }

        isClosingSession = false
        hasEstablishedSession = true
        scanner.delegate = self
        if shouldScanWhenReady {
            shouldScanWhenReady = false
            configureFunctionalUnit(on: scanner)
            state = .scanning
            scanner.requestScan()
        } else {
            state = .ready
        }
    }

    private func processSessionClosed(on device: ICDevice, error: Error?) {
        let wasClosing = isClosingSession
        resumeSessionCloseWait()

        guard let scanner = device as? ICScannerDevice,
              selectedDevice === scanner else { return }

        if wasClosing {
            if state != .idle { state = .browsing }
            return
        }

        // Many scanners close the session after each scan — reopen quietly for the next one.
        guard state != .idle else { return }
        maintainSession(for: scanner)
    }

    /// Keeps the scanner session open after each scan for the next one.
    private func resetAfterScan() {
        clearConnectionTimeout()
        scanCompletion = nil
        scannedImageURL = nil
        didFinishCurrentScan = false
        accumulatedBands = []
        shouldScanWhenReady = false
        isOpeningSession = false

        if let device = selectedDevice {
            device.delegate = self
            if device.hasOpenSession {
                state = .ready
            } else {
                state = .ready
                maintainSession(for: device)
            }
        }
    }

    private var connectionTimeoutTask: Task<Void, Never>?

    private func scheduleConnectionTimeout(for device: ICScannerDevice) {
        connectionTimeoutTask?.cancel()
        connectionTimeoutTask = Task { @MainActor in
            try? await Task.sleep(nanoseconds: connectionTimeoutSeconds * 1_000_000_000)
            guard !Task.isCancelled else { return }
            guard selectedDevice === device, case .connecting = state else { return }
            isOpeningSession = false
            device.delegate = nil
            state = .error("Could not connect to the scanner. Select it again to retry.")
        }
    }

    private func clearConnectionTimeout() {
        connectionTimeoutTask?.cancel()
        connectionTimeoutTask = nil
    }

    private func configureFunctionalUnit(on device: ICScannerDevice) {
        let unit = device.selectedFunctionalUnit
        unit.pixelDataType = .RGB
        if let eightBit = ICScannerBitDepth(rawValue: 8) {
            unit.bitDepth = eightBit
        }
        unit.resolution = resolution.dpi

        let docType = documentType(for: paperSize)
        if let flatbed = unit as? ICScannerFunctionalUnitFlatbed {
            flatbed.documentType = docType
        } else if let feeder = unit as? ICScannerFunctionalUnitDocumentFeeder {
            feeder.documentType = docType
        } else if let positive = unit as? ICScannerFunctionalUnitPositiveTransparency {
            positive.documentType = docType
        } else if let negative = unit as? ICScannerFunctionalUnitNegativeTransparency {
            negative.documentType = docType
        }

        let physicalSize = unit.physicalSize
        if physicalSize.width > 0, physicalSize.height > 0 {
            unit.scanArea = NSRect(origin: .zero, size: physicalSize)
        }

        device.transferMode = .fileBased
        device.documentName = "Scan"
        device.documentUTI = "public.jpeg"
        device.downloadsDirectory = scanDownloadsDirectory()
    }

    private func scanDownloadsDirectory() -> URL {
        let directory = FileManager.default.temporaryDirectory
            .appendingPathComponent("Scans", isDirectory: true)
        try? FileManager.default.createDirectory(at: directory, withIntermediateDirectories: true)
        return directory
    }

    private func documentType(for paperSize: PaperSize) -> ICScannerDocumentType {
        switch paperSize {
        case .a4: return .typeA4
        case .letter: return .typeUSLetter
        case .legal: return .typeUSLegal
        }
    }

    private func isSharedDevice(_ device: ICDevice) -> Bool {
        let locationBits = device.type.rawValue & ICDeviceLocationTypeMask.remote.rawValue
        return locationBits != 0
    }

    private func categorize(_ device: ICScannerDevice, moreComing: Bool) {
        if isSharedDevice(device) {
            if !sharedDevices.contains(where: { $0 === device }) {
                sharedDevices.append(device)
            }
        } else if !localDevices.contains(where: { $0 === device }) {
            localDevices.append(device)
        }
        autoSelectIfNeeded(device, moreComing: moreComing)
    }

    private func autoSelectIfNeeded(_ device: ICScannerDevice, moreComing: Bool) {
        guard selectedDevice == nil else { return }

        let deviceName = device.name ?? device.locationDescription
        if let preferred = pendingPreferredDeviceName,
           let deviceName,
           deviceName == preferred {
            pendingPreferredDeviceName = nil
            selectDevice(device)
            return
        }

        if !moreComing, localDevices.count + sharedDevices.count == 1 {
            selectDevice(device)
        }
    }

    private func remove(_ device: ICDevice) {
        localDevices.removeAll { $0 === device }
        sharedDevices.removeAll { $0 === device }
        if selectedDevice === device {
            selectedDevice = nil
            state = .browsing
        }
    }

    private func finishScan(with result: Result<NSImage, Error>) {
        guard !didFinishCurrentScan else { return }
        didFinishCurrentScan = true
        scanCompletion?(result)
        resetAfterScan()
    }

    private func imageFromScanURL(_ url: URL) -> NSImage? {
        if let data = try? Data(contentsOf: url), let image = NSImage(data: data) {
            return image
        }
        return NSImage(contentsOf: url)
    }

    private func loadImageFromScanURL(_ url: URL) async -> NSImage? {
        for attempt in 0..<6 {
            if let image = imageFromScanURL(url) {
                return image
            }
            if attempt < 5 {
                try? await Task.sleep(nanoseconds: 250_000_000)
            }
        }
        return nil
    }

    private func resolveScannedImage() async -> NSImage? {
        if let url = scannedImageURL, let image = await loadImageFromScanURL(url) {
            return image
        }
        if !accumulatedBands.isEmpty {
            return assembleImage(from: accumulatedBands)
        }
        return nil
    }
}

// MARK: - ICDeviceBrowserDelegate

extension ScannerService: ICDeviceBrowserDelegate {
    nonisolated func deviceBrowser(_ browser: ICDeviceBrowser, didAdd device: ICDevice, moreComing: Bool) {
        guard let scanner = device as? ICScannerDevice else { return }
        Task { @MainActor in
            categorize(scanner, moreComing: moreComing)
        }
    }

    nonisolated func deviceBrowser(_ browser: ICDeviceBrowser, didRemove device: ICDevice, moreGoing: Bool) {
        Task { @MainActor in
            remove(device)
        }
    }
}

// MARK: - ICDeviceDelegate / ICScannerDeviceDelegate

extension ScannerService: ICDeviceDelegate, ICScannerDeviceDelegate {
    nonisolated func didRemove(_ device: ICDevice) {
        Task { @MainActor in
            remove(device)
        }
    }

    nonisolated func device(_ device: ICDevice, didOpenSessionWithError error: (any Error)?) {
        Task { @MainActor in
            guard let scanner = device as? ICScannerDevice else { return }
            processSessionOpened(on: scanner, error: error)
        }
    }

    nonisolated func device(_ device: ICDevice, didCloseSessionWithError error: (any Error)?) {
        Task { @MainActor in
            processSessionClosed(on: device, error: error)
        }
    }

    nonisolated func scannerDevice(_ scanner: ICScannerDevice, didScanTo url: URL) {
        Task { @MainActor in
            scannedImageURL = url
            if let image = imageFromScanURL(url) {
                finishScan(with: .success(image))
            }
        }
    }

    nonisolated func scannerDevice(_ scanner: ICScannerDevice, didScanTo bandData: ICScannerBandData) {
        Task { @MainActor in
            accumulatedBands.append(bandData)
            if bandData.dataNumRows >= bandData.fullImageHeight,
               let image = imageFromBandData(bandData) {
                finishScan(with: .success(image))
            }
        }
    }

    nonisolated func scannerDevice(_ scanner: ICScannerDevice, didCompleteScanWithError error: (any Error)?) {
        Task { @MainActor in
            if didFinishCurrentScan { return }
            if let error {
                finishScan(with: .failure(error))
                return
            }
            if let image = await resolveScannedImage() {
                finishScan(with: .success(image))
                return
            }
            finishScan(with: .failure(ScanError.scanFailed))
        }
    }

    nonisolated func scannerDeviceDidBecomeAvailable(_ scanner: ICScannerDevice) {
        Task { @MainActor in
            guard selectedDevice === scanner else { return }
            guard state == .connecting || isOpeningSession else { return }
            clearConnectionTimeout()
            isOpeningSession = false
            isClosingSession = false
            hasEstablishedSession = true
            scanner.delegate = self
            if shouldScanWhenReady {
                shouldScanWhenReady = false
                configureFunctionalUnit(on: scanner)
                state = .scanning
                scanner.requestScan()
            } else {
                state = .ready
            }
        }
    }
}

private extension ScanResolution {
    var dpi: Int {
        switch self {
        case .dpi150: return 150
        case .dpi300: return 300
        case .dpi600: return 600
        }
    }
}

private extension ScannerService {
    func imageFromBandData(_ bandData: ICScannerBandData) -> NSImage? {
        guard let data = bandData.dataBuffer,
              !data.isEmpty,
              bandData.fullImageWidth > 0,
              bandData.dataNumRows > 0 else { return nil }

        let colorSpace = CGColorSpaceCreateDeviceRGB()
        let bitmapInfo = CGBitmapInfo(rawValue: CGImageAlphaInfo.noneSkipLast.rawValue)
        guard let provider = CGDataProvider(data: data as CFData),
              let cgImage = CGImage(
                  width: Int(bandData.fullImageWidth),
                  height: Int(bandData.dataNumRows),
                  bitsPerComponent: Int(bandData.bitsPerComponent),
                  bitsPerPixel: Int(bandData.bitsPerPixel),
                  bytesPerRow: Int(bandData.bytesPerRow),
                  space: colorSpace,
                  bitmapInfo: bitmapInfo,
                  provider: provider,
                  decode: nil,
                  shouldInterpolate: true,
                  intent: .defaultIntent
              ) else {
            return nil
        }

        let rep = NSBitmapImageRep(cgImage: cgImage)
        let image = NSImage(size: NSSize(width: bandData.fullImageWidth, height: bandData.dataNumRows))
        image.addRepresentation(rep)
        return verticallyFlipped(image)
    }

    func assembleImage(from bands: [ICScannerBandData]) -> NSImage? {
        guard let first = bands.first,
              first.fullImageWidth > 0,
              first.fullImageHeight > 0 else { return nil }

        let width = Int(first.fullImageWidth)
        let height = Int(first.fullImageHeight)
        let bytesPerRow = Int(first.bytesPerRow)
        guard bytesPerRow > 0 else { return nil }

        var fullData = Data(count: bytesPerRow * height)
        for band in bands {
            guard let bandData = band.dataBuffer, !bandData.isEmpty else { continue }
            let startRow = Int(band.dataStartRow)
            let startOffset = startRow * bytesPerRow
            let copyLength = min(bandData.count, fullData.count - startOffset)
            guard startOffset >= 0, copyLength > 0, startOffset + copyLength <= fullData.count else { continue }
            fullData.replaceSubrange(startOffset..<(startOffset + copyLength), with: bandData.prefix(copyLength))
        }

        let colorSpace = CGColorSpaceCreateDeviceRGB()
        let bitmapInfo = CGBitmapInfo(rawValue: CGImageAlphaInfo.noneSkipLast.rawValue)
        var mutableData = fullData
        return mutableData.withUnsafeMutableBytes { buffer in
            guard let baseAddress = buffer.baseAddress,
                  let context = CGContext(
                      data: baseAddress,
                      width: width,
                      height: height,
                      bitsPerComponent: Int(first.bitsPerComponent),
                      bytesPerRow: bytesPerRow,
                      space: colorSpace,
                      bitmapInfo: bitmapInfo.rawValue
                  ),
                  let cgImage = context.makeImage() else {
                return nil
            }
            let rep = NSBitmapImageRep(cgImage: cgImage)
            let image = NSImage(size: NSSize(width: width, height: height))
            image.addRepresentation(rep)
            return verticallyFlipped(image)
        }
    }

    func verticallyFlipped(_ image: NSImage) -> NSImage {
        guard let cgImage = image.cgImage(forProposedRect: nil, context: nil, hints: nil) else { return image }
        let width = cgImage.width
        let height = cgImage.height
        let colorSpace = cgImage.colorSpace ?? CGColorSpaceCreateDeviceRGB()
        guard let context = CGContext(
                  data: nil,
                  width: width,
                  height: height,
                  bitsPerComponent: cgImage.bitsPerComponent,
                  bytesPerRow: 0,
                  space: colorSpace,
                  bitmapInfo: cgImage.bitmapInfo.rawValue
              ) else {
            return image
        }
        context.translateBy(x: 0, y: CGFloat(height))
        context.scaleBy(x: 1, y: -1)
        context.draw(cgImage, in: CGRect(x: 0, y: 0, width: width, height: height))
        guard let flipped = context.makeImage() else { return image }
        return NSImage(cgImage: flipped, size: NSSize(width: width, height: height))
    }
}