stamp.swift

#!/usr/bin/swift

import AppKit

// MARK: - Helpers

@inline(__always) func error(_ message: String) {
    print("💥 \(message)")
}
@inline(__always) func success(_ message: String) {
    print("🎉 \(message)")
}

extension NSImage {
    
    var bitmapRepresentation: NSBitmapImageRep? {
        guard let image = cgImage(forProposedRect: nil, context: nil, hints: nil) else { return nil }
        return NSBitmapImageRep(cgImage: image)
    }
}

extension String {
    
    func finished(with end: String) -> String {
        guard hasSuffix(end) == false else { return self }
        return self + end
    }
}

// MARK: - Attribute

enum Attribute {
    case font(NSFont)
    case shadow(color: NSColor, offset: NSSize, blur: CGFloat)
    case alignment(NSTextAlignment)
    case lineBreakMode(NSLineBreakMode)
    case color(NSColor)
    
    var type: String {
        switch self {
        case .font:
            return "font"
        case .shadow:
            return "shadow"
        case .alignment:
            return "alignment"
        case .lineBreakMode:
            return "lineBreakMode"
        case .color:
            return "color"
        }
    }
}

extension Attribute: Hashable {
    
    static func ==(lhs: Attribute, rhs: Attribute) -> Bool {
        return lhs.type == rhs.type
    }
    
    var hashValue: Int {
        return type.hashValue
    }
}

// MARK: - Styler

struct Styler {
    
    private(set) var attributes: Set<Attribute>
    
    init(attributes: Attribute...) {
        self.init(attributes: attributes)
    }
    
    init(attributes: [Attribute] = []) {
        self.attributes = Set(attributes)
    }
    
    mutating func add(_ attribute: Attribute) {
        attributes.insert(attribute)
    }
    
    var dictionary: [String: AnyObject] {
        
        var result: [String: AnyObject] = [:]
        
        let paragraphStyle = NSParagraphStyle.default().mutableCopy() as! NSMutableParagraphStyle
        
        for attribute in attributes {
            switch attribute {
            case let .font(font):
                result[NSFontAttributeName] = font
            case let .shadow(color, offset, blur):
                let shadow = NSShadow()
                shadow.shadowColor = color
                shadow.shadowOffset = offset
                shadow.shadowBlurRadius = blur
                result[NSShadowAttributeName] = shadow
            case let .alignment(alignment):
                paragraphStyle.alignment = alignment
            case let .lineBreakMode(mode):
                paragraphStyle.lineBreakMode = mode
            case let .color(color):
                result[NSForegroundColorAttributeName] = color
            }
        }
        
        result[NSParagraphStyleAttributeName] = paragraphStyle
        
        return result
    }
}

// MARK: - Stamper

final class Stamper {
    
    private struct Constants {
        private init() {}
        
        static let scale: CGFloat = {
            NSScreen.screens()?.flatMap { $0.backingScaleFactor }.max() ?? 1
        }()
        static let fontName = "HelveticaNeue-Light"
        static let minimumFontSize: CGFloat = 6
    }
    
    private var icon: NSImage?
    private var path: String?
    
    func load(path: String) -> Stamper {
        
        guard let icon = NSImage(contentsOfFile: path) else {
            error("Couldn't load image at: \(path)")
            return self
        }
        
        self.icon = icon
        self.path = path
        
        return self
    }
    
    func overlay(with layer: NSImage?) -> Stamper {
        
        guard let layer = layer, let icon = icon else { return self }
        
        let image = layer.copy() as! NSImage
        let result = icon.copy() as! NSImage
        
        result.lockFocus()
        image.size = result.size
        image.draw(at: .zero, from: NSRect.init(origin: .zero, size: result.size), operation: .plusDarker, fraction: 1.0)
        result.unlockFocus()
        
        self.icon = result
        
        return self
    }
    
    func stamp(text: String, styler: Styler) -> Stamper {
        
        guard let icon = icon else { return self }
        
        let originalSize = icon.size
        let scale = Constants.scale
        
        icon.size = NSSize(width: originalSize.width / scale,
                           height: originalSize.height / scale)
        
        let offset = icon.size.height / 20
        let containerSize = CGSize(width: icon.size.width,
                                   height: icon.size.height - 2 * offset)
        var fontSize = icon.size.height / 4
        
        var textContainer: NSTextContainer
        var textStorage: NSTextStorage
        var layoutManager: NSLayoutManager
        
        var renderedRange: NSRange
        var usedRect: CGRect
        
        var styler = styler
        
        let unit = icon.size.height / 64
        styler.add(.shadow(color: .controlDarkShadowColor, offset: NSSize(width: 0, height: -unit), blur: unit))
        
        repeat {
            
            if let font = NSFont(name: Constants.fontName, size: fontSize) {
                styler.add(.font(font))
            }
            
            textContainer = NSTextContainer(containerSize: containerSize)
            textStorage = NSTextStorage(string: text, attributes: styler.dictionary)
            
            layoutManager = NSLayoutManager()
            layoutManager.addTextContainer(textContainer)
            layoutManager.textStorage = textStorage
            
            renderedRange = layoutManager.glyphRange(for: textContainer)
            usedRect = layoutManager.usedRect(for: textContainer)
            
            fontSize -= 0.1
            
        } while renderedRange.length < text.characters.count && fontSize > Constants.minimumFontSize
        
        let point = CGPoint(x: 0, y: icon.size.height - usedRect.size.height - offset)
        
        icon.lockFocusFlipped(true)
        layoutManager.drawGlyphs(forGlyphRange: renderedRange, at: point)
        icon.unlockFocus()
        
        icon.size = originalSize
        
        return self
    }
    
    func save(path: String? = nil) {
        
        guard let path = path ?? self.path else { return }
        let url = URL(fileURLWithPath: path)
        
        do {
            try icon?
                .bitmapRepresentation?
                .representation(using: .PNG, properties: [:])?
                .write(to: url, options: .atomic)
        } catch _ {
            error("Couldn't save image at: \(path)")
        }
    }
}

// MARK: - Crawler

final class Crawler {
    
    func lookup(name: String, at path: String) -> [URL] {
        
        let baseURL = URL(fileURLWithPath: path)
        
        return FileManager.default
            .enumerator(atPath: path)?
            .flatMap { $0 as? String }
            .filter { $0.hasSuffix(name) }
            .flatMap { baseURL.appendingPathComponent($0) } ?? []
    }
}

// MARK: - Worker

final class Worker {
    
    private var iconSets: [URL] = []
    
    func find(iconSet: String, at path: String) -> Worker {
        
        iconSets = Crawler().lookup(name: iconSet, at: path)
        return self
    }
    
    func stamp(tags: [String], styler: Styler, overlay: NSImage?) {
        
        for set in iconSets {
            for tag in tags {
                stamp(tag: tag, on: set, styler: styler, overlay: overlay)
            }
        }
    }
    
    private func stamp(tag: String, on iconSet: URL, styler: Styler, overlay: NSImage?) {
        
        let taggedIconSet = self.url(for: iconSet, with: tag)
        
        if FileManager.default.fileExists(atPath: taggedIconSet.path) {
            do {
                try FileManager.default.removeItem(at: taggedIconSet)
            } catch _ {
                error("Couldn't remove existing tagged icon set at: \(taggedIconSet.path)")
            }
        }
        
        do {
            try FileManager.default.copyItem(at: iconSet, to: taggedIconSet)
        } catch _ {
            error("Could create a copy of the icon set at: \(iconSet.path)")
        }
        
        Crawler()
            .lookup(name: ".png", at: taggedIconSet.path)
            .forEach { stamp(icon: $0, with: tag, styler: styler, overlay: overlay) }
    }
    
    private func stamp(icon: URL, with tag: String, styler: Styler, overlay: NSImage?) {
        
        Stamper()
            .load(path: icon.path)
            .overlay(with: overlay)
            .stamp(text: tag, styler: styler)
            .save()
    }
    
    private func url(for iconSet: URL, with tag: String) -> URL {
        
        let name = iconSet.deletingPathExtension().lastPathComponent + tag.capitalized
        let url = iconSet
            .deletingLastPathComponent()
            .appendingPathComponent(name)
            .appendingPathExtension(iconSet.pathExtension)
        
        return url
    }
}

// MARK: - Main

guard 3...4 ~= CommandLine.arguments.count else {
    print("Usage: stamp.swift PATH APP_ICON_SET [OVERLAY]")
    exit(1)
}

var path = CommandLine.arguments[1]
let iconSet = CommandLine.arguments[2].finished(with: ".appiconset")
let overlay = CommandLine.arguments.count == 4
    ? NSImage(contentsOfFile: CommandLine.arguments[3])
    : nil

let tags: [String] = ["dev", "test", "preprod"]
let styler = Styler(attributes: .alignment(.center),
                    .lineBreakMode(.byWordWrapping),
                    .color(.white))

Worker()
    .find(iconSet: iconSet, at: path)
    .stamp(tags: tags, styler: styler, overlay: overlay)

success("Done")