smart_printer/smart_printer/OCRFileView.swift

import Cocoa
import UniformTypeIdentifiers
import Vision

// MARK: - OCR Service

enum OCRError: LocalizedError {
    case invalidImage
    case recognitionFailed
    case noTextFound

    var errorDescription: String? {
        switch self {
        case .invalidImage:
            return "The selected file could not be opened as an image."
        case .recognitionFailed:
            return "Text recognition failed. Please try another image."
        case .noTextFound:
            return "No text was found in this image."
        }
    }
}

enum OCRFileService {
    private static let photoExtensions: Set<String> = [
        "jpg", "jpeg", "png", "heic", "heif", "gif", "tiff", "tif", "bmp", "webp",
    ]

    static func present(from window: NSWindow?) {
        let hostWindow = window ?? NSApp.keyWindow
        guard PremiumAccess.require(feature: .ocrFile, from: hostWindow) else { return }

        let panel = NSOpenPanel()
        panel.title = "OCR File"
        panel.message = "Select an image to extract text from."
        panel.prompt = "Select Image"
        panel.canChooseDirectories = false
        panel.canChooseFiles = true
        panel.allowsMultipleSelection = false
        panel.allowsOtherFileTypes = false
        panel.allowedContentTypes = photoExtensions.compactMap { UTType(filenameExtension: $0) }
        panel.directoryURL = FileManager.default.urls(for: .picturesDirectory, in: .userDomainMask).first

        let completion: (NSApplication.ModalResponse) -> Void = { response in
            guard response == .OK, let url = panel.url else { return }
            let hostWindow = window ?? NSApp.keyWindow
            guard let viewController = PremiumAccess.mainViewController(from: hostWindow) else { return }
            viewController.presentOCRFile(imageURL: url)
        }

        if let window = window ?? NSApp.keyWindow {
            panel.beginSheetModal(for: window, completionHandler: completion)
        } else if panel.runModal() == .OK {
            completion(.OK)
        }
    }

    static func recognizeText(from url: URL) async throws -> String {
        let accessed = url.startAccessingSecurityScopedResource()
        defer {
            if accessed { url.stopAccessingSecurityScopedResource() }
        }

        guard let image = NSImage(contentsOf: url),
              let cgImage = image.cgImage(forProposedRect: nil, context: nil, hints: nil) else {
            throw OCRError.invalidImage
        }

        return try await recognizeText(in: cgImage)
    }

    static func recognizeText(from image: NSImage) async throws -> String {
        guard let cgImage = image.cgImage(forProposedRect: nil, context: nil, hints: nil) else {
            throw OCRError.invalidImage
        }
        return try await recognizeText(in: cgImage)
    }

    /// Recognizes text from each image independently. Returns an empty string for pages with no text.
    static func recognizeTextPages(from images: [NSImage]) async -> [String] {
        var results: [String] = []
        for image in images {
            if let text = try? await recognizeText(from: image) {
                results.append(text.trimmingCharacters(in: .whitespacesAndNewlines))
            } else {
                results.append("")
            }
        }
        return results
    }

    /// Recognizes text from multiple scan pages and joins them with blank lines between pages.
    static func recognizeText(from images: [NSImage]) async -> String {
        let pageTexts = await recognizeTextPages(from: images).filter { !$0.isEmpty }
        return pageTexts.joined(separator: "\n\n")
    }

    static func recognizeText(in cgImage: CGImage) async throws -> String {
        let observations = try await recognizeTextObservations(in: cgImage)
        let text = textFromObservations(observations)
        guard !text.isEmpty else { throw OCRError.noTextFound }
        return text
    }

    static func recognizeTextObservations(in cgImage: CGImage) async throws -> [VNRecognizedTextObservation] {
        try await withCheckedThrowingContinuation { continuation in
            let request = VNRecognizeTextRequest { request, error in
                if let error {
                    continuation.resume(throwing: error)
                    return
                }
                guard let observations = request.results as? [VNRecognizedTextObservation] else {
                    continuation.resume(throwing: OCRError.recognitionFailed)
                    return
                }
                continuation.resume(returning: observations)
            }
            request.recognitionLevel = .accurate
            request.usesLanguageCorrection = true

            let handler = VNImageRequestHandler(cgImage: cgImage, options: [:])
            do {
                try handler.perform([request])
            } catch {
                continuation.resume(throwing: error)
            }
        }
    }

    static func textFromObservations(_ observations: [VNRecognizedTextObservation]) -> String {
        let sorted = observations.sorted { lhs, rhs in
            let leftY = 1 - lhs.boundingBox.midY
            let rightY = 1 - rhs.boundingBox.midY
            if abs(leftY - rightY) > 0.02 { return leftY < rightY }
            return lhs.boundingBox.minX < rhs.boundingBox.minX
        }
        return sorted
            .compactMap { $0.topCandidates(1).first?.string }
            .joined(separator: "\n")
            .trimmingCharacters(in: .whitespacesAndNewlines)
    }
}

// MARK: - Overlay

final class OCRFileOverlayView: NSView, AppearanceRefreshable {
    var onDismiss: (() -> Void)?

    private let imageURL: URL
    private let backButton = OCRFileToolbarButton(symbolName: "chevron.left", accessibilityLabel: "Back")
    private let titleLabel = NSTextField(labelWithString: "OCR File")
    private let freeUsesBanner = FreeUsesBannerView()
    private let textContainer = NSView()
    private let textScrollView = NSScrollView()
    private let textView = NSTextView()
    private let loadingIndicator = NSProgressIndicator()
    private let statusLabel = NSTextField(labelWithString: "Extracting text…")
    private let printButton = OCRFilePrintButton()

    init(imageURL: URL) {
        self.imageURL = imageURL
        super.init(frame: .zero)
        translatesAutoresizingMaskIntoConstraints = false
        setup()
        refreshAppearance()
        NotificationCenter.default.addObserver(
            self,
            selector: #selector(appearanceDidChange),
            name: .appearanceDidChange,
            object: nil
        )
        recognizeText()
    }

    deinit {
        NotificationCenter.default.removeObserver(self)
    }

    @objc private func appearanceDidChange() {
        refreshAppearance()
    }

    @available(*, unavailable)
    required init?(coder: NSCoder) { nil }

    func refreshAppearance() {
        layer?.backgroundColor = AppTheme.background.cgColor
        titleLabel.textColor = AppTheme.textPrimary
        textContainer.layer?.backgroundColor = AppTheme.cardBackground.cgColor
        textContainer.layer?.borderColor = AppTheme.paywallBorder.cgColor
        statusLabel.textColor = AppTheme.textSecondary
        backButton.refreshAppearance()
        printButton.refreshAppearance()
        freeUsesBanner.refreshAppearance()
    }

    func present(in parent: NSView) {
        guard superview == nil else { return }
        parent.addSubview(self)
        NSLayoutConstraint.activate([
            leadingAnchor.constraint(equalTo: parent.leadingAnchor),
            trailingAnchor.constraint(equalTo: parent.trailingAnchor),
            topAnchor.constraint(equalTo: parent.topAnchor),
            bottomAnchor.constraint(equalTo: parent.bottomAnchor),
        ])
        alphaValue = 0
        NSAnimationContext.runAnimationGroup { context in
            context.duration = 0.2
            animator().alphaValue = 1
        }
    }

    func dismiss(animated: Bool = true) {
        let remove = { [weak self] in
            self?.removeFromSuperview()
            self?.onDismiss?()
        }
        guard animated else {
            remove()
            return
        }
        NSAnimationContext.runAnimationGroup({ context in
            context.duration = 0.15
            animator().alphaValue = 0
        }, completionHandler: remove)
    }

    private func setup() {
        wantsLayer = true

        titleLabel.font = AppTheme.semiboldFont(size: 18)
        titleLabel.alignment = .center
        titleLabel.translatesAutoresizingMaskIntoConstraints = false

        textContainer.wantsLayer = true
        textContainer.layer?.cornerRadius = AppTheme.contentPanelCornerRadius
        textContainer.layer?.borderWidth = 1.5
        textContainer.applyCardShadow()
        textContainer.translatesAutoresizingMaskIntoConstraints = false

        textScrollView.translatesAutoresizingMaskIntoConstraints = false
        textScrollView.hasVerticalScroller = true
        textScrollView.hasHorizontalScroller = false
        textScrollView.autohidesScrollers = true
        textScrollView.borderType = .noBorder
        textScrollView.drawsBackground = false

        textView.isRichText = false
        textView.isEditable = true
        textView.isSelectable = true
        textView.drawsBackground = false
        textView.textColor = AppTheme.textPrimary
        textView.font = AppTheme.regularFont(size: 14)
        textView.textContainerInset = NSSize(width: 16, height: 16)
        textView.isAutomaticQuoteSubstitutionEnabled = true
        textView.isAutomaticDashSubstitutionEnabled = true
        textView.isAutomaticTextReplacementEnabled = true
        textView.isHidden = true

        textScrollView.documentView = textView

        loadingIndicator.style = .spinning
        loadingIndicator.controlSize = .regular
        loadingIndicator.translatesAutoresizingMaskIntoConstraints = false
        loadingIndicator.startAnimation(nil)

        statusLabel.font = AppTheme.regularFont(size: 14)
        statusLabel.alignment = .center
        statusLabel.translatesAutoresizingMaskIntoConstraints = false

        backButton.translatesAutoresizingMaskIntoConstraints = false
        printButton.translatesAutoresizingMaskIntoConstraints = false
        printButton.isEnabled = false

        backButton.onClick = { [weak self] in self?.dismiss() }
        printButton.onClick = { [weak self] in self?.printText() }

        addSubview(backButton)
        addSubview(titleLabel)
        addSubview(freeUsesBanner)
        addSubview(textContainer)
        textContainer.addSubview(textScrollView)
        textContainer.addSubview(loadingIndicator)
        textContainer.addSubview(statusLabel)
        addSubview(printButton)

        NSLayoutConstraint.activate([
            backButton.leadingAnchor.constraint(equalTo: leadingAnchor, constant: 24),
            backButton.topAnchor.constraint(equalTo: topAnchor, constant: 20),
            backButton.widthAnchor.constraint(equalToConstant: 40),
            backButton.heightAnchor.constraint(equalToConstant: 40),

            titleLabel.centerXAnchor.constraint(equalTo: centerXAnchor),
            titleLabel.centerYAnchor.constraint(equalTo: backButton.centerYAnchor),

            freeUsesBanner.leadingAnchor.constraint(equalTo: leadingAnchor, constant: AppTheme.contentPanelInset),
            freeUsesBanner.topAnchor.constraint(equalTo: backButton.bottomAnchor, constant: 16),

            textContainer.leadingAnchor.constraint(equalTo: leadingAnchor, constant: AppTheme.contentPanelInset),
            textContainer.trailingAnchor.constraint(equalTo: trailingAnchor, constant: -AppTheme.contentPanelInset),
            textContainer.topAnchor.constraint(equalTo: freeUsesBanner.bottomAnchor, constant: 16),
            textContainer.bottomAnchor.constraint(equalTo: printButton.topAnchor, constant: -28),

            textScrollView.leadingAnchor.constraint(equalTo: textContainer.leadingAnchor, constant: 12),
            textScrollView.trailingAnchor.constraint(equalTo: textContainer.trailingAnchor, constant: -12),
            textScrollView.topAnchor.constraint(equalTo: textContainer.topAnchor, constant: 12),
            textScrollView.bottomAnchor.constraint(equalTo: textContainer.bottomAnchor, constant: -12),

            loadingIndicator.centerXAnchor.constraint(equalTo: textContainer.centerXAnchor),
            loadingIndicator.centerYAnchor.constraint(equalTo: textContainer.centerYAnchor, constant: -12),

            statusLabel.topAnchor.constraint(equalTo: loadingIndicator.bottomAnchor, constant: 12),
            statusLabel.leadingAnchor.constraint(equalTo: textContainer.leadingAnchor, constant: 24),
            statusLabel.trailingAnchor.constraint(equalTo: textContainer.trailingAnchor, constant: -24),

            printButton.centerXAnchor.constraint(equalTo: centerXAnchor),
            printButton.bottomAnchor.constraint(equalTo: bottomAnchor, constant: -28),
            printButton.widthAnchor.constraint(greaterThanOrEqualToConstant: 220),
        ])
    }

    private func recognizeText() {
        Task { @MainActor in
            do {
                let text = try await OCRFileService.recognizeText(from: imageURL)
                showRecognizedText(text)
            } catch {
                showRecognitionError(error)
            }
        }
    }

    private func showRecognizedText(_ text: String) {
        loadingIndicator.stopAnimation(nil)
        loadingIndicator.isHidden = true
        statusLabel.isHidden = true
        textView.string = text
        textView.isHidden = false
        printButton.isEnabled = true
        window?.makeFirstResponder(textView)
    }

    private func showRecognitionError(_ error: Error) {
        loadingIndicator.stopAnimation(nil)
        loadingIndicator.isHidden = true
        statusLabel.isHidden = true

        let alert = NSAlert()
        alert.messageText = "OCR Failed"
        alert.informativeText = error.localizedDescription
        alert.alertStyle = .warning
        alert.addButton(withTitle: "OK")
        alert.beginSheetModal(for: window ?? NSApp.keyWindow ?? NSWindow()) { [weak self] _ in
            self?.dismiss()
        }
    }

    private func printText() {
        PrintService.printText(textView.string, title: "OCR File", from: window)
    }
}

// MARK: - Toolbar Button

private final class OCRFileToolbarButton: NSControl, AppearanceRefreshable {
    var onClick: (() -> Void)?

    private let symbolName: String
    private let iconView = NSImageView()

    init(symbolName: String, accessibilityLabel: String) {
        self.symbolName = symbolName
        super.init(frame: .zero)
        toolTip = accessibilityLabel

        wantsLayer = true
        layer?.cornerRadius = 12

        if let image = NSImage(systemSymbolName: symbolName, accessibilityDescription: accessibilityLabel) {
            let config = NSImage.SymbolConfiguration(pointSize: 15, weight: .semibold)
            iconView.image = image.withSymbolConfiguration(config)
        }
        iconView.translatesAutoresizingMaskIntoConstraints = false

        addSubview(iconView)
        NSLayoutConstraint.activate([
            iconView.centerXAnchor.constraint(equalTo: centerXAnchor),
            iconView.centerYAnchor.constraint(equalTo: centerYAnchor),
        ])
        refreshAppearance()
    }

    @available(*, unavailable)
    required init?(coder: NSCoder) { nil }

    func refreshAppearance() {
        layer?.backgroundColor = AppTheme.elevatedBackground.cgColor
        layer?.borderColor = AppTheme.paywallBorder.cgColor
        layer?.borderWidth = 1.5
        iconView.contentTintColor = AppTheme.textPrimary
    }

    override func mouseUp(with event: NSEvent) {
        guard bounds.contains(convert(event.locationInWindow, from: nil)) else { return }
        onClick?()
    }

    override func resetCursorRects() {
        addCursorRect(bounds, cursor: .pointingHand)
    }
}

// MARK: - Print Button

private final class OCRFilePrintButton: NSControl, AppearanceRefreshable {
    var onClick: (() -> Void)?

    private let titleLabel = NSTextField(labelWithString: "Print Text")
    private let iconView = NSImageView()
    private var hoverTracker: HoverTracker?
    private var isHovered = false
    private var isEnabledState = true

    override var isEnabled: Bool {
        get { isEnabledState }
        set {
            isEnabledState = newValue
            alphaValue = newValue ? 1 : 0.45
        }
    }

    init() {
        super.init(frame: .zero)
        wantsLayer = true
        layer?.cornerRadius = 22

        titleLabel.font = AppTheme.semiboldFont(size: 16)
        titleLabel.textColor = .white
        titleLabel.translatesAutoresizingMaskIntoConstraints = false

        if let image = NSImage(systemSymbolName: "printer.fill", accessibilityDescription: "Print") {
            let config = NSImage.SymbolConfiguration(pointSize: 14, weight: .semibold)
            iconView.image = image.withSymbolConfiguration(config)
        }
        iconView.contentTintColor = .white
        iconView.translatesAutoresizingMaskIntoConstraints = false

        addSubview(titleLabel)
        addSubview(iconView)

        NSLayoutConstraint.activate([
            heightAnchor.constraint(equalToConstant: 52),

            titleLabel.leadingAnchor.constraint(equalTo: leadingAnchor, constant: 28),
            titleLabel.centerYAnchor.constraint(equalTo: centerYAnchor),

            iconView.leadingAnchor.constraint(equalTo: titleLabel.trailingAnchor, constant: 10),
            iconView.centerYAnchor.constraint(equalTo: centerYAnchor),
            iconView.trailingAnchor.constraint(equalTo: trailingAnchor, constant: -28),
            iconView.widthAnchor.constraint(equalToConstant: 20),
            iconView.heightAnchor.constraint(equalToConstant: 20),
        ])

        hoverTracker = HoverTracker(view: self) { [weak self] hovering in
            self?.setHovered(hovering)
        }
        refreshAppearance()
    }

    @available(*, unavailable)
    required init?(coder: NSCoder) { nil }

    func refreshAppearance() {
        let color = isHovered
            ? AppTheme.blue.blended(withFraction: 0.15, of: .black) ?? AppTheme.blue
            : AppTheme.blue
        layer?.backgroundColor = color.cgColor
        titleLabel.textColor = .white
        iconView.contentTintColor = .white
    }

    private func setHovered(_ hovering: Bool) {
        guard isEnabledState else { return }
        isHovered = hovering
        animateHover {
            let color = hovering
                ? AppTheme.blue.blended(withFraction: 0.15, of: .black) ?? AppTheme.blue
                : AppTheme.blue
            layer?.backgroundColor = color.cgColor
            layer?.transform = hovering
                ? CATransform3DMakeScale(1.03, 1.03, 1)
                : CATransform3DIdentity
        }
    }

    override func mouseUp(with event: NSEvent) {
        guard isEnabledState else { return }
        guard bounds.contains(convert(event.locationInWindow, from: nil)) else { return }
        onClick?()
    }

    override func resetCursorRects() {
        guard isEnabledState else { return }
        addCursorRect(bounds, cursor: .pointingHand)
    }
}