HTMLStreamer/Sources/HTMLStreamer/AttributedStringConverter.swift

//
//  AttributedStringConverter.swift
//  HTMLStreamer
//
//  Created by Shadowfacts on 11/24/23.
//

#if os(iOS)
import UIKit
#elseif os(macOS)
import AppKit
#endif

struct AttributedStringConverter<Callbacks: AttributedStringCallbacks> {
    private let configuration: AttributedStringConverterConfiguration
    private var tokenizer: Tokenizer<String.Iterator>
    private let str = NSMutableAttributedString()
    
    private var actionStack: InlineArray3<ElementAction> = []
    private var styleStack: InlineArray3<Style> = []
    // The current run of text w/o styles changing
    private var currentRun: String = ""
    
    init(html: String, configuration: AttributedStringConverterConfiguration) where Callbacks == DefaultCallbacks {
        self.init(html: html, configuration: configuration, callbacks: DefaultCallbacks.self)
    }
    
    init(html: String, configuration: AttributedStringConverterConfiguration, callbacks _: Callbacks.Type = Callbacks.self) {
        self.configuration = configuration
        self.tokenizer = Tokenizer(chars: html.makeIterator())
    }
    
    mutating func convert() -> NSAttributedString {
        while let token = tokenizer.next() {
            switch token {
            case .character(let c):
                currentRun.append(c)
            case .comment:
                // ignored
                continue
            case .startTag(let name, let selfClosing, let attributes):
                let action = Callbacks.elementAction(name: name, attributes: attributes)
                actionStack.append(action)
                // self closing tags are ignored since they have no content
                if !selfClosing {
                    handleStartTag(name, attributes: attributes)
                }
            case .endTag(let name):
                handleEndTag(name)
                // if we have a non-default action for the current element, the run finishes here
                if actionStack.last != .default {
                    finishRun()
                }
                actionStack.removeLast()
            case .doctype:
                // ignored
                continue
            }
        }
        
        finishRun()
        
        return str
    }
    
    private mutating func handleStartTag(_ name: String, attributes: InlineArray3<HTMLStreamer.Attribute>) {
        switch name {
        case "br":
            currentRun.append("\n")
        case "a":
            // we need to always insert in attribute, because we need to always have one
            // to remove from the stack in handleEndTag
            // but we only need to finish the run if we have a URL, since otherwise
            // the final attribute run won't be affected
            let url = attributes.attributeValue(for: "href").flatMap(Callbacks.makeURL(string:))
            if url != nil {
                finishRun()
            }
            styleStack.append(.link(url))
        case "em", "i":
            finishRun()
            styleStack.append(.italic)
        case "strong", "b":
            finishRun()
            styleStack.append(.bold)
        case "del":
            finishRun()
            styleStack.append(.strikethrough)
        case "code":
            finishRun()
            styleStack.append(.monospace)
        case "pre":
            startBlockElement()
            finishRun()
            styleStack.append(.monospace)
        case "blockquote":
            startBlockElement()
            finishRun()
            styleStack.append(.blockquote)
        case "p":
            startBlockElement()
        default:
            break
        }
    }
    
    private mutating func startBlockElement() {
        if str.length != 0 || !currentRun.isEmpty {
            currentRun.append("\n\n")
        }
    }
    
    private mutating func handleEndTag(_ name: String) {
        switch name {
        case "a":
            if case .link(.some(_)) = styleStack.last {
                finishRun()
            }
            removeLastStyle(.link)
        case "em", "i":
            finishRun()
            removeLastStyle(.italic)
        case "strong", "b":
            finishRun()
            removeLastStyle(.bold)
        case "del":
            finishRun()
            removeLastStyle(.strikethrough)
        case "code":
            finishRun()
            removeLastStyle(.monospace)
        case "pre":
            finishRun()
            removeLastStyle(.monospace)
        case "blockquote":
            finishRun()
            removeLastStyle(.blockquote)
        default:
            break
        }
    }
    
    // needed to correctly handle mis-nested tags
    private mutating func removeLastStyle(_ type: Style.StyleType) {
        var i = styleStack.index(before: styleStack.endIndex)
        while i >= styleStack.startIndex {
            if styleStack[i].type == type {
                styleStack.remove(at: i)
                return
            }
            styleStack.formIndex(before: &i)
        }
    }
    
    private lazy var blockquoteParagraphStyle: NSParagraphStyle = {
        let style = configuration.paragraphStyle.mutableCopy() as! NSMutableParagraphStyle
        style.headIndent = 32
        style.firstLineHeadIndent = 32
        return style
    }()
    
    private mutating func finishRun() {
        guard !currentRun.isEmpty else {
            return
        }
        
        if actionStack.contains(.skip) {
            currentRun = ""
            return
        } else if case .replace(let replacement) = actionStack.first(where: \.isReplace) {
            currentRun = replacement
        }
        
        var attributes = [NSAttributedString.Key: Any]()
        var currentFontTraits = Set<FontTrait>()
        for style in styleStack {
            switch style {
            case .bold:
                currentFontTraits.insert(.bold)
            case .italic:
                currentFontTraits.insert(.italic)
            case .monospace:
                currentFontTraits.insert(.monospace)
            case .link(let url):
                if let url {
                    attributes[.link] = url
                }
            case .strikethrough:
                attributes[.strikethroughStyle] = NSUnderlineStyle.single.rawValue
            case .blockquote:
                attributes[.paragraphStyle] = blockquoteParagraphStyle
                currentFontTraits.insert(.italic)
            }
        }
        
        let baseFont = currentFontTraits.contains(.monospace) ? configuration.monospaceFont : configuration.font
        var descriptor = baseFont.fontDescriptor
        if currentFontTraits.contains(.bold) && currentFontTraits.contains(.italic),
           let boldItalic = descriptor.withSymbolicTraits([.traitBold, .traitItalic]) {
            descriptor = boldItalic
        } else if currentFontTraits.contains(.bold),
           let bold = descriptor.withSymbolicTraits(.traitBold) {
            descriptor = bold
        } else if currentFontTraits.contains(.italic),
           let italic = descriptor.withSymbolicTraits(.traitItalic) {
            descriptor = italic
        }
        #if os(iOS)
        attributes[.font] = UIFont(descriptor: descriptor, size: 0)
        #elseif os(macOS)
        attributes[.font] = NSFont(descriptor: descriptor, size: 0)
        #endif
        
        if !attributes.keys.contains(.paragraphStyle) {
            attributes[.paragraphStyle] = configuration.paragraphStyle
        }
        
        str.append(NSAttributedString(string: currentRun, attributes: attributes))
        currentRun = ""
    }
}

protocol AttributedStringCallbacks {
    static func makeURL(string: String) -> URL?
    static func elementAction(name: String, attributes: InlineArray3<Attribute>) -> ElementAction
}

enum ElementAction: Equatable {
    case `default`
    case skip
    case replace(String)
    
    var isReplace: Bool {
        if case .replace(_) = self {
            true
        } else {
            false
        }
    }
}

extension AttributedStringCallbacks {
    static func makeURL(string: String) -> URL? {
        URL(string: string)
    }
    static func elementAction(name: String, attributes: InlineArray3<Attribute>) -> ElementAction {
        .default
    }
}

struct DefaultCallbacks: AttributedStringCallbacks {
}

struct AttributedStringConverterConfiguration {
    #if os(iOS)
    var font: UIFont
    var monospaceFont: UIFont
    var color: UIColor
    #elseif os(macOS)
    var font: NSFont
    var monospaceFont: NSFont
    var color: NSColor
    #endif
    var paragraphStyle: NSParagraphStyle
}

#if os(macOS)
private extension NSFontDescriptor {
    func withSymbolicTraits(_ traits: SymbolicTraits) -> NSFontDescriptor? {
        let descriptor: NSFontDescriptor = self.withSymbolicTraits(traits)
        return descriptor
    }
}
private extension NSFontDescriptor.SymbolicTraits {
    static var traitBold: Self { .bold }
    static var traitItalic: Self { .italic }
}
#endif

private enum FontTrait {
    case bold
    case italic
    case monospace
}

private enum Style {
    case bold
    case italic
    case monospace
    case link(URL?)
    case strikethrough
    case blockquote
    
    var type: StyleType {
        switch self {
        case .bold:
            return .bold
        case .italic:
            return .italic
        case .monospace:
            return .monospace
        case .link(_):
            return .link
        case .strikethrough:
            return .strikethrough
        case .blockquote:
            return .blockquote
        }
    }
    
    enum StyleType {
        case bold, italic, monospace, link, strikethrough, blockquote
    }
}

extension Collection where Element == Attribute {
    func attributeValue(for name: String) -> String? {
        first(where: { $0.name == name })?.value
    }
}