#!/usr/bin/swift
//add a caption to an image file
// Jerry Stratton astoundingscripts.com

import Quartz

//standard values
let validFiletypes = [".jpg", ".tiff", ".pdf", ".png"]
let files = FileManager.default
let alignments = [
	"left": NSTextAlignment.left,
	"center": NSTextAlignment.center,
	"right": NSTextAlignment.right,
	"justify": NSTextAlignment.justified,
]
let cases = ["lower", "title", "upper"]

//command-line arguments
var arguments = CommandLine.arguments
var commandName = arguments.removeFirst()
func help(message: String = "") {
	print("Syntax:", commandName, "<image> [caption] [caption 2-x] [--options]")
	print("\timage: path to the image to caption")
	print("\tcaption: caption text to add to bottom of image; use 'none' to have no text")
	print("\t--help: print this help and exit")
	print("\t--align <" + alignments.keys.joined(separator:"|") + ">: force a specific alignment")
	print("\t--bgcolor <r,g,b,[f]|FFFFFF[FF]>: background color for the caption text")
	print("\t--border [percentage]: add border around image relative to image width")
	print("\t--case <" + cases.joined(separator:"|") + ">: transform text case")
	print("\t--fgcolor <r,g,b,[f]|FFFFFF[FF]>: foreground color for the caption text")
	print("\t--font [name] [size] [bold|italic]>: font name and optionally size, bold, and/or italic")
	print("\t--layer [auto|margin]: layer caption over the image instead of above or below; auto: create a left bug; margin: 0-100")
	print("\t--output <path>: image to create; must end in one of:", validFiletypes.joined(separator:", "))
	print("\t--padding <multiplier>: increase or decrease padding around the caption")
	print("\t--quote: add quotes to caption if it doesn’t already have them")
	print("\t--signature <text>: add signature to lower right")
	print("\t--top: put the margin on the top instead of the bottom")
	print("\t--width <width>: scale final image to requested pixel width")
	if message != "" {
		print("\n" + message + "\n")
	}
	exit(0)
}

//image manipulation routines
func scaleImage(image:NSImage, width:CGFloat) -> NSImage {
	//calculate new height from width
	let height = image.size.height*width/image.size.width
	let scaledSize = NSSize(width: Int(width), height: Int(height))

	//resize image
	let scaledImage = NSImage(size:scaledSize, flipped: false) { (resizedRect) -> Bool in
		image.draw(in: resizedRect)
		return true
	}

	return scaledImage
}

//return a font with bold and/or italics
func getStyledFont(font:NSFont, bold:Bool=false, italic:Bool=false) -> NSFont {
	var traits = font.fontDescriptor.symbolicTraits
	if bold {
		traits.insert(.bold)
	}
	if italic {
		traits.insert(.italic)
	}
	//this apears to produce a styled version even if none exists
	//although it won't necessarily produce *both* styles if two exist
	let descriptor = font.fontDescriptor.withFamily(font.familyName!).withSymbolicTraits(traits)
	let styledFont = NSFont(descriptor: descriptor, size: font.pointSize)
	if styledFont != nil {
		return styledFont!
	} else {
		print("Cannot get styled version of font. Using unstyled version.");
		return font;
	}
}

//parse arguments
func popFloat(warning: String = "", defaultValue: CGFloat? = nil) -> CGFloat {
	var value:Float? = nil
	if arguments.count > 0 {
		value = Float(arguments[0])
	}
	if value == nil {
		if defaultValue != nil {
			return(defaultValue!)
		} else {
			help(message:warning)
		}
	} else {
		arguments.removeFirst()
	}
	return CGFloat(value!)
}

func popString(warning: String) -> String {
	var value:String = ""
	if arguments.count > 0 {
		value = arguments.removeFirst()
	} else {
		help(message:warning)
	}
	return value
}

extension String {
	func isFloat() -> Bool {
		if Float(self) != nil {
			return true
		}
		return false
	}

	func isFont() -> Bool {
		guard nil != NSFont(name: self, size: 12) else {
			return false
		}
		return true
	}
}

func stringToColor(color:String) -> NSColor {
	var rgbf = [CGFloat]()
	if color.count >= 4 && color.count % 2 == 0 && color.range(of: "^[0-9A-Fa-f]+$", options: .regularExpression) != nil {
		//hex color
		var hexString = color
		while hexString.count > 0 {
			let hex = String(hexString.removeFirst()) + String(hexString.removeFirst())
			guard let hexNumber = UInt8(hex, radix:16) else {
				print(hex, "is not a valid hexadecimal number.")
				exit(0)
			}
			rgbf.append(CGFloat(hexNumber)/255)
		}
	} else if color.count >= 5 && color.range(of: "^[0-9.,]+$", options: .regularExpression) != nil {
		//comma-delimited floating point color
		let parts = color.components(separatedBy: ",")
		rgbf = parts.map { CGFloat(($0 as NSString).floatValue) }
	} else if color.range(of: "^[0-9.]+$", options: .regularExpression) != nil {
		//greyscale floating point color
		rgbf = [CGFloat](repeating: CGFloat(Float(color) ?? 0), count: 3)
	} else {
		print("String", color, "not recognized as color format.")
		exit(0)
	}

	while rgbf.count < 4 {
		rgbf.append(CGFloat(1.0))
	}
	return NSColor(red:rgbf[0], green:rgbf[1], blue:rgbf[2], alpha:rgbf[3])
}

var text = ""
var alignment = ""
var imageFile:NSImage? = nil
var backgroundColor = NSColor(red:1, green:1, blue:1, alpha:1)
var bold = false
var border:CGFloat = 0.0
var caseTransformation = "none"
var colorSet = false
var enquote = false
var fontName = "Adobe Garamond"
var fontSize:CGFloat = 24.0
var foregroundColor = NSColor(red:0, green:0, blue:0, alpha:1)
var italic = false
var layerCaption = false
var layerAutoWidth = false
var margin:CGFloat = 0.0
var padding:CGFloat = 1.0
var outputFile = ""
var outputWidth:CGFloat = 0.0
var signature = ""
var topCaption = false

while arguments.count > 0 {
	let argument = arguments.removeFirst()
	switch argument {
		case "--help", "-h":
			help()
		case "--align":
			alignment = popString(warning:"--align requires a value")
			if !alignments.keys.contains(alignment) {
				help(message:"Text alignment “" + alignment + "” not recognized.")
			}
		case "--bgcolor", "--fgcolor":
			let color = popString(warning:"--bgcolor and --fgcolor must be color values.")
			colorSet = true
			if argument == "--bgcolor" {
				backgroundColor = stringToColor(color:color)
			} else {
				foregroundColor = stringToColor(color:color)
			}
		case "--border", "-b":
			border = popFloat(defaultValue:0.67)
		case "--case":
			caseTransformation = popString(warning:"--case requires one of " + cases.joined(separator:", "))
			if !cases.contains(caseTransformation) {
				help(message:"Case transformation must be one of " + cases.joined(separator:", "))
			}
		case "--font", "-f":
			var newFontSize:CGFloat = 0.0
			var newFontName = ""
			var hadFontOption = false
			fontOptions:while arguments.count > 0 {
				switch arguments[0] {
					case "bold":
						bold = true
					case "italic":
						italic = true
					case let option where option.hasPrefix("--"):
						break fontOptions
					default:
						//might be a font size
						if newFontSize == 0.0 && arguments[0].isFloat() {
							newFontSize = CGFloat(Float(arguments[0]) ?? 0)
							if newFontSize != 0 {
								fontSize = newFontSize
								break
							}
						}

						//might be a valid font name
						if newFontName == "" && arguments[0].isFont() {
							newFontName = arguments[0]
							fontName = newFontName
							break
						}

						//if this isn't a valid font option, and there have been no other options
						//assume this is a misspelled font option
						if !hadFontOption {
							help(message:newFontName + " is not a valid font name or option")
						}

						//no longer in font options
						break fontOptions
				}
				arguments.removeFirst()
				hadFontOption = true
			}
			if !hadFontOption {
				help(message:"--font requires at least one option")
			}
		case "--layer", "-l":
			layerCaption = true
			while arguments.count > 0 {
				if arguments[0] == "auto" {
					arguments.removeFirst()
					layerAutoWidth = true
				} else if arguments[0].isFloat() {
					var helpMessage = "margins must be between 0 and 100"
					margin = popFloat(warning:helpMessage)
					if margin <= 0 || margin >= 100 {
						if margin == 0 {
							helpMessage += "\nzero is unnecessary; it is the default layer location"
						} else if margin == 100 {
							helpMessage += "\nuse --top to move the layered caption to the top"
						}
						help(message:helpMessage)
					}
				} else {
					break
				}
			}
		case "--output", "-o":
			outputFile = popString(warning:"--output requires a filepath.")
			if !validFiletypes.contains(where:outputFile.hasSuffix) {
				help(message:"Output image must end in: " + validFiletypes.joined(separator:", "))
			}
		case "--padding", "-p":
			padding = popFloat(warning:"--padding requires a positive number")
			if padding <= 0 {
				help(message:"padding must be positive")
			}
		case "--quote":
			enquote = true
		case "--signature":
			signature = popString(warning: "--signature requires signature text")
		case "--top", "-t":
			topCaption = true
		case "--width":
			outputWidth = popFloat(warning:"--width requires a value")
		default:
			if imageFile == nil && files.fileExists(atPath:argument) {
				imageFile = NSImage(byReferencingFile: argument)
				if !imageFile!.isValid {
					print(argument, "is not a readable image.")
					exit(0)
				}
				if outputFile == "" {
					var filename = (argument as NSString).lastPathComponent
					filename = (filename as NSString).deletingPathExtension
					outputFile = filename + " captioned.jpg"
				}
				break
			}

			if !argument.starts(with:"--") {
				if text == "" {
					text = argument
				} else {
					text += "\n" + argument
				}
				break
			}
			help(message:"Unknown option: " + argument)
	}
}

//image is required
if imageFile == nil {
	help(message:"An image file must be specified.")
}

//read text if there is none
if text == "none" {
	text = ""
} else if text == "" {
	while let line = readLine() {
		if text != "" {
			text += "\n"
		}
		text += line
	}
}
switch caseTransformation {
	case "lower":
		text = text.localizedLowercase
		signature = signature.localizedLowercase
	case "title":
		text = text.localizedCapitalized
	case "upper":
		text = text.localizedUppercase
		signature = signature.localizedUppercase
	default:
		break
}

//draw the image and the text
//outputSize is the size of the overall captioned image file
var outputSize = imageFile!.size
//imageRect is the size of the actual image being captioned
//if you wanted to crop the image, imageRect is what you would modify
var imageRect = NSRect(x:0, y:0, width:outputSize.width, height:outputSize.height)
//imageArea is where the image will be drawn in the captioned image
var imageArea = NSRect(x:0, y:0, width:imageRect.width, height:imageRect.height)
var borderColor = foregroundColor

//create the caption box and add to the size
var caption = NSAttributedString()
//captionRect is the entire caption area with its padding
var captionRect = NSRect()
//captionArea is where the text is drawn
var captionArea = NSRect()
var captionAttributes: [NSAttributedString.Key: Any] = [:]
if text != "" {
	//swap background and foreground if layering and still using the default colors
	if layerCaption && !colorSet {
		swap(&foregroundColor, &backgroundColor)
		backgroundColor = backgroundColor.withAlphaComponent(0.67)
	}
	captionAttributes[NSAttributedString.Key.foregroundColor] = foregroundColor
	
	let captionStyle = NSMutableParagraphStyle()
	captionAttributes[NSAttributedString.Key.paragraphStyle] = captionStyle
	var alignmentGuess = NSTextAlignment.center

	//arbitrarily make the font size relative to the width of the image
	fontSize *= imageRect.width/800
	var font = NSFont(name: fontName, size: fontSize)!

	//bold and italics
	if bold || italic {
		font = getStyledFont(font:font, bold:bold, italic:italic)
	}

	//quote it if we want quotes and it isn't quoted already
	if enquote && !text.hasPrefix("“") {
		text = "“" + text + "”"
	}

	//create the text for drawing
	captionAttributes[NSAttributedString.Key.font] = font
	caption = NSAttributedString(string:text, attributes:captionAttributes)

	//spacing around the top and bottom of the caption
	let singleLine = NSAttributedString(string:"etaonri", attributes:captionAttributes).size()
	let characterWidth = singleLine.width/7
	let captionPadding = characterWidth*padding
	//captionAreaSize is the width and height of the actual text, without padding
	var captionAreaSize = NSSize(width:outputSize.width - captionPadding*2, height:0)

	//if this is more than one line, it needs to be justified instead of centered
	var captionSize = caption.boundingRect(with:captionAreaSize, options:NSString.DrawingOptions.usesLineFragmentOrigin)
	if captionSize.width > captionAreaSize.width*0.8 && captionSize.height >= singleLine.height*2 {
		alignmentGuess = NSTextAlignment.justified
		let quoteCharacter = NSAttributedString(string:"“", attributes:captionAttributes)
		captionStyle.firstLineHeadIndent = characterWidth
		captionStyle.headIndent = characterWidth
		if text.first == "“" {
			captionStyle.headIndent += quoteCharacter.size().width
		}
		//redo the captionSize in case changing the alignment modifies it
		captionSize = caption.boundingRect(with:captionAreaSize, options:NSString.DrawingOptions.usesLineFragmentOrigin)
	}

	//create the caption area, and adjust the output size and image area
	//captionRectSize is the width and height of the caption with padding
	var captionRectSize = NSSize(width:outputSize.width, height:captionSize.height+captionPadding*2)
	if layerCaption {
		if layerAutoWidth {
			captionRectSize.width = captionSize.width + captionPadding*2
			captionAreaSize.width = captionSize.width
		}
	} else {
		outputSize.height += captionSize.height + captionPadding*2
		if !topCaption {
			imageArea = NSOffsetRect(imageArea, 0, captionSize.height+captionPadding*2)
		}
	}
	let captionLeft = outputSize.width-captionRectSize.width
	var captionFloor:CGFloat = 0.0
	if topCaption {
		captionFloor = outputSize.height - captionSize.height - captionPadding*2
		if margin > 0 {
			captionFloor -= imageRect.height*margin/100
		}
	} else if margin > 0 {
		captionFloor += imageRect.height*margin/100
	}
	captionRect = NSRect(x:captionLeft, y:captionFloor, width:captionRectSize.width, height:captionRectSize.height)
	captionArea = NSRect(x:captionLeft+captionPadding, y:captionFloor+captionPadding, width:captionAreaSize.width, height:captionSize.height)

	//set the alignment to the requested or to the guessed
	captionStyle.alignment = alignments[alignment] ?? alignmentGuess
}

//add the border to the size
if border > 0 {
	border = imageRect.width*border/100
	outputSize.height += border*2
	outputSize.width += border*2
	imageArea = NSOffsetRect(imageArea, border, border)
	captionArea = NSOffsetRect(captionArea, border, border)
	captionRect = NSOffsetRect(captionRect, border, border)
}

//add signature
//create signature text based off of main text
func signatureMaker(string:String, fontRatio:CGFloat=1) -> NSAttributedString {
	var signatureAttributes = captionAttributes
	if (fontRatio != 1.0) {
		let signatureFont = NSFont(name:fontName, size:fontSize/fontRatio)
		signatureAttributes[NSAttributedString.Key.font] = signatureFont
	}
	return NSAttributedString(string:signature, attributes:signatureAttributes)
}
var signatureText = signatureMaker(string:signature)
var signatureRect = NSRect()
if signature != "" {
	var signatureSize = signatureText.size()

	//the signature should be no more than a third the width of the image
	if signatureSize.width > outputSize.width/3 {
		let signatureRatio:CGFloat = signatureSize.width/(outputSize.width/3)
		signatureText = signatureMaker(string:signature, fontRatio:signatureRatio)
		signatureSize = signatureText.size()
	}

	let signatureX = outputSize.width*0.995-signatureSize.width-border
	let signatureY = outputSize.width*0.005+border
	signatureRect = NSRect(x:signatureX, y:signatureY, width:signatureSize.width, height:signatureSize.height)
}

//create the captioned image
var outputImage = NSImage(size:outputSize, flipped: false) { (outputRect) -> Bool in
	imageFile!.draw(in:imageArea, from:imageRect, operation:NSCompositingOperation.copy, fraction:1.0)

	//add the caption, if any
	if text != "" {
		backgroundColor.setFill()
		captionRect.fill()
		caption.draw(in: captionArea)
	}

	//add the signature if requested
	if signature != "" {
		signatureText.draw(in:signatureRect)
	}

	//draw the border if requested
	if border > 0 {
		borderColor.setFill()
		outputRect.frame(withWidth:border)
	}

	return true
}

if (outputWidth > 0.0) {
	outputImage = scaleImage(image:outputImage, width:outputWidth)
}

//convert image to desired format
var imageData:Data?
if outputFile.hasSuffix(".pdf") {
	guard let pdfPage = PDFPage(image: outputImage) else {
		print("Unable to acquire PDF data.")
		exit(0)
	}
	imageData = pdfPage.dataRepresentation
} else {
	imageData = outputImage.tiffRepresentation
	if !outputFile.hasSuffix(".tiff") {
		let bitmapVersion = NSBitmapImageRep(data: imageData!)
		var filetype = NSBitmapImageRep.FileType.png
		if outputFile.hasSuffix(".jpg") {
			filetype = NSBitmapImageRep.FileType.jpeg
		}
		imageData = bitmapVersion!.representation(using: filetype, properties:[:])
	}
}

//write image to file
do {
	try imageData!.write(to: NSURL.fileURL(withPath: outputFile))
} catch {
	print("Unable to save file: ", outputFile)
}
