Snaggle/Snaggle/Editor/Surface.swift

//  Surface.swift
//  Snaggle
//
//  Port of Greenshot's Surface: the image plus its drawable elements, undo and rendering.

import AppKit

final class Surface {

    private(set) var baseImage: CGImage
    /// Size of the image in points; elements live in this coordinate space.
    private(set) var imageSize: CGSize
    private(set) var elements: [Element] = []

    var details: CaptureDetails
    var selection: Element?
    var cropElement: CropElement?
    var nextCounterNumber: Int = 1
    var hasUnsavedChanges = false

    let undoManager = UndoManager()
    var onChange: (() -> Void)?

    init(capture: Capture) {
        self.baseImage = capture.image
        self.imageSize = capture.pointSize
        self.details = capture.details
    }

    var scale: CGFloat { CGFloat(baseImage.width) / max(1, imageSize.width) }

    // MARK: - Element management

    func add(_ element: Element) {
        elements.append(element)
        if let crop = element as? CropElement { cropElement = crop }
        undoManager.registerUndo(withTarget: self) { surface in
            surface.remove(element)
        }
        undoManager.setActionName("Add \(String(describing: type(of: element)))")
        changed()
    }

    func remove(_ element: Element) {
        guard let index = elements.firstIndex(where: { $0 === element }) else { return }
        elements.remove(at: index)
        if element === cropElement { cropElement = nil }
        if element === selection { selection = nil }
        undoManager.registerUndo(withTarget: self) { surface in
            surface.add(element)
        }
        changed()
    }

    func removeSelection() {
        guard let selection else { return }
        remove(selection)
    }

    func bringToFront(_ element: Element) {
        guard let index = elements.firstIndex(where: { $0 === element }) else { return }
        elements.append(elements.remove(at: index))
        changed()
    }

    func sendToBack(_ element: Element) {
        guard let index = elements.firstIndex(where: { $0 === element }) else { return }
        elements.insert(elements.remove(at: index), at: 0)
        changed()
    }

    func duplicate(_ element: Element) {
        let copy = element.copyElement()
        copy.move(by: CGVector(dx: 12, dy: 12))
        add(copy)
        selection = copy
    }

    /// Records the state of an element before an interactive change so undo can restore it.
    func registerGeometryUndo(for element: Element) {
        let previousFrame = element.frame
        undoManager.registerUndo(withTarget: self) { surface in
            surface.registerGeometryUndo(for: element)
            element.frame = previousFrame
            surface.changed()
        }
    }

    func element(at point: CGPoint) -> Element? {
        elements.reversed().first { !($0 is CropElement) && $0.hitTest(point) }
    }

    func changed() {
        hasUnsavedChanges = true
        onChange?()
    }

    // MARK: - Crop

    func applyCrop(to rect: CGRect) {
        let clamped = rect.intersection(CGRect(origin: .zero, size: imageSize))
        guard clamped.width >= 1, clamped.height >= 1 else { return }

        let previousImage = baseImage
        let previousSize = imageSize
        let previousElements = elements

        let pixelRect = CGRect(x: (clamped.minX * scale).rounded(.down),
                               y: (clamped.minY * scale).rounded(.down),
                               width: (clamped.width * scale).rounded(),
                               height: (clamped.height * scale).rounded())
        guard let cropped = baseImage.cropping(to: pixelRect) else { return }

        undoManager.registerUndo(withTarget: self) { surface in
            surface.restore(image: previousImage, size: previousSize, elements: previousElements)
        }
        undoManager.setActionName("Crop")

        baseImage = cropped
        imageSize = clamped.size
        elements.removeAll { $0 is CropElement }
        cropElement = nil
        for element in elements {
            element.move(by: CGVector(dx: -clamped.minX, dy: -clamped.minY))
        }
        selection = nil
        changed()
    }

    private func restore(image: CGImage, size: CGSize, elements restored: [Element]) {
        let currentImage = baseImage
        let currentSize = imageSize
        let currentElements = elements
        undoManager.registerUndo(withTarget: self) { surface in
            surface.restore(image: currentImage, size: currentSize, elements: currentElements)
        }
        baseImage = image
        imageSize = size
        elements = restored
        cropElement = nil
        selection = nil
        changed()
    }

    // MARK: - Rendering

    /// Renders the base image with every element baked in, at full pixel resolution.
    func render() -> CGImage {
        let width = baseImage.width
        let height = baseImage.height
        guard width > 0, height > 0,
              let context = CGContext(data: nil,
                                      width: width,
                                      height: height,
                                      bitsPerComponent: 8,
                                      bytesPerRow: 0,
                                      space: CGColorSpace(name: CGColorSpace.sRGB) ?? CGColorSpaceCreateDeviceRGB(),
                                      bitmapInfo: CGImageAlphaInfo.premultipliedLast.rawValue)
        else { return baseImage }

        // Flip to a top-left origin and work in points.
        context.translateBy(x: 0, y: CGFloat(height))
        context.scaleBy(x: 1, y: -1)
        context.scaleBy(x: scale, y: scale)
        context.interpolationQuality = .high

        drawFlipped(baseImage, in: CGRect(origin: .zero, size: imageSize), context: context)

        for element in elements where !(element is CropElement) {
            element.draw(in: context, baseImage: baseImage, imageSize: imageSize)
        }
        return context.makeImage() ?? baseImage
    }

    func renderedCapture() -> Capture {
        let capture = Capture(image: render(), details: details)
        capture.details.scaleFactor = scale
        return capture
    }
}