Sapling/Sources/SaplingEditor/HybridMarkdownLineRenderer.swift

import Foundation

public enum HybridMarkdownLineKind: Hashable, Sendable {
    case paragraph
    case heading(level: Int, markerRange: NSRange, textRange: NSRange)
    case blockquote(markerRange: NSRange, contentRange: NSRange)
    case horizontalRule(range: NSRange)
    case unorderedList(markerRange: NSRange, contentRange: NSRange, nestingLevel: Int)
    case orderedList(markerRange: NSRange, contentRange: NSRange, nestingLevel: Int)
    case taskList(
        markerRange: NSRange,
        checkboxRange: NSRange,
        contentRange: NSRange,
        checked: Bool,
        nestingLevel: Int
    )
    case fencedCodeFence(markerRange: NSRange, languageRange: NSRange?)
    case codeBlockContent
    case tableRow(cellRanges: [NSRange], separatorRanges: [NSRange], isDivider: Bool)
}

public enum HybridMarkdownSpanKind: Hashable, Sendable {
    case bold
    case italic
    case inlineCode
    case link
    case automaticLink
    case markdownDelimiter
}

public struct HybridMarkdownSpan: Hashable, Sendable {
    public var range: NSRange
    public var kind: HybridMarkdownSpanKind

    public init(range: NSRange, kind: HybridMarkdownSpanKind) {
        self.range = range
        self.kind = kind
    }
}

public struct HybridMarkdownLineRenderPlan: Hashable, Sendable {
    public var line: EditorLine
    public var kind: HybridMarkdownLineKind
    public var spans: [HybridMarkdownSpan]

    public init(line: EditorLine, kind: HybridMarkdownLineKind, spans: [HybridMarkdownSpan]) {
        self.line = line
        self.kind = kind
        self.spans = spans
    }
}

public struct HybridMarkdownLineRenderer: Sendable {
    public init() {}

    public func renderPlan(for line: EditorLine) -> HybridMarkdownLineRenderPlan {
        let kind = lineKind(for: line)
        let spans = inlineSpans(in: line, kind: kind)
        return HybridMarkdownLineRenderPlan(line: line, kind: kind, spans: spans)
    }

    public func renderPlans(for lines: [EditorLine]) -> [HybridMarkdownLineRenderPlan] {
        var isInCodeBlock = false
        var plans: [HybridMarkdownLineRenderPlan] = []

        for line in lines {
            let nonCodeKind = lineKind(for: line)
            let kind: HybridMarkdownLineKind

            if case .fencedCodeFence = nonCodeKind {
                kind = nonCodeKind
                isInCodeBlock.toggle()
            } else if isInCodeBlock {
                kind = .codeBlockContent
            } else {
                kind = nonCodeKind
            }

            plans.append(HybridMarkdownLineRenderPlan(
                line: line,
                kind: kind,
                spans: inlineSpans(in: line, kind: kind)
            ))
        }

        return plans
    }

    private func lineKind(for line: EditorLine) -> HybridMarkdownLineKind {
        let firstToken = firstNonWhitespaceCharacter(in: line.source)

        if (firstToken == "`" || firstToken == "~"), let fence = fencedCodeFence(in: line) {
            return fence
        }

        if (firstToken == "-" || firstToken == "*" || firstToken == "_"), isHorizontalRule(line.source) {
            return .horizontalRule(range: line.range)
        }

        if firstToken == "#", let heading = heading(in: line) {
            return heading
        }

        if (firstToken == "-" || firstToken == "*"), let task = taskList(in: line) {
            return task
        }

        if (firstToken == "-" || firstToken == "*"), let unorderedList = unorderedList(in: line) {
            return unorderedList
        }

        if firstToken?.isNumber == true, let orderedList = orderedList(in: line) {
            return orderedList
        }

        if firstToken == ">", let blockquote = blockquote(in: line) {
            return blockquote
        }

        if let table = tableRow(in: line) {
            return table
        }

        return .paragraph
    }

    private func heading(in line: EditorLine) -> HybridMarkdownLineKind? {
        let leadingWhitespace = line.source.prefix { $0 == " " }.count
        guard leadingWhitespace <= 3 else { return nil }

        let headingStart = line.source.dropFirst(leadingWhitespace)
        let markerCount = headingStart.prefix { $0 == "#" }.count
        guard (1...6).contains(markerCount),
              headingStart.dropFirst(markerCount).first == " "
        else {
            return nil
        }

        let textOffset = leadingWhitespace + markerCount + 1
        return .heading(
            level: markerCount,
            markerRange: NSRange(location: line.range.location + leadingWhitespace, length: markerCount),
            textRange: NSRange(
                location: line.range.location + textOffset,
                length: max(0, line.range.length - textOffset)
            )
        )
    }

    private func blockquote(in line: EditorLine) -> HybridMarkdownLineKind? {
        let trimmedPrefix = line.source.prefix { $0 == " " }
        guard trimmedPrefix.count <= 3 else { return nil }
        let start = trimmedPrefix.count
        guard line.source.dropFirst(start).first == ">" else { return nil }

        var markerLength = start
        var cursor = line.source.index(line.source.startIndex, offsetBy: start)
        while cursor < line.source.endIndex, line.source[cursor] == ">" {
            markerLength += 1
            cursor = line.source.index(after: cursor)
        }
        if cursor < line.source.endIndex, line.source[cursor] == " " {
            markerLength += 1
        }

        return .blockquote(
            markerRange: NSRange(location: line.range.location, length: markerLength),
            contentRange: NSRange(
                location: line.range.location + markerLength,
                length: max(0, line.range.length - markerLength)
            )
        )
    }

    private func unorderedList(in line: EditorLine) -> HybridMarkdownLineKind? {
        guard let match = firstMatch(Patterns.unorderedList, in: line.source),
              match.range(at: 3).length > 0
        else { return nil }

        let indentation = match.range(at: 1).length
        let markerRange = NSRange(
            location: line.range.location + match.range(at: 2).location,
            length: match.range(at: 2).length
        )
        let contentRange = NSRange(
            location: line.range.location + match.range(at: 3).location,
            length: match.range(at: 3).length
        )
        return .unorderedList(markerRange: markerRange, contentRange: contentRange, nestingLevel: nestingLevel(indentation))
    }

    private func orderedList(in line: EditorLine) -> HybridMarkdownLineKind? {
        guard let match = firstMatch(Patterns.orderedList, in: line.source),
              match.range(at: 3).length > 0
        else { return nil }

        let indentation = match.range(at: 1).length
        let markerRange = NSRange(
            location: line.range.location + match.range(at: 2).location,
            length: match.range(at: 2).length
        )
        let contentRange = NSRange(
            location: line.range.location + match.range(at: 3).location,
            length: match.range(at: 3).length
        )
        return .orderedList(markerRange: markerRange, contentRange: contentRange, nestingLevel: nestingLevel(indentation))
    }

    private func taskList(in line: EditorLine) -> HybridMarkdownLineKind? {
        guard let match = firstMatch(Patterns.taskList, in: line.source),
              match.range(at: 4).length > 0
        else { return nil }

        let indentation = match.range(at: 1).length
        let checkbox = (line.source as NSString).substring(with: match.range(at: 3))
        return .taskList(
            markerRange: NSRange(
                location: line.range.location + match.range(at: 2).location,
                length: match.range(at: 2).length
            ),
            checkboxRange: NSRange(
                location: line.range.location + match.range(at: 3).location,
                length: match.range(at: 3).length
            ),
            contentRange: NSRange(
                location: line.range.location + match.range(at: 4).location,
                length: match.range(at: 4).length
            ),
            checked: checkbox.lowercased() == "[x]",
            nestingLevel: nestingLevel(indentation)
        )
    }

    private func fencedCodeFence(in line: EditorLine) -> HybridMarkdownLineKind? {
        guard let match = firstMatch(Patterns.fencedCodeFence, in: line.source) else { return nil }
        let languageRange = match.range(at: 3).length > 0
            ? NSRange(location: line.range.location + match.range(at: 3).location, length: match.range(at: 3).length)
            : nil
        return .fencedCodeFence(
            markerRange: NSRange(
                location: line.range.location + match.range(at: 2).location,
                length: match.range(at: 2).length
            ),
            languageRange: languageRange
        )
    }

    private func tableRow(in line: EditorLine) -> HybridMarkdownLineKind? {
        guard line.source.contains("|") else { return nil }
        let trimmed = line.source.trimmingCharacters(in: .whitespaces)
        guard trimmed.contains("|"), trimmed.count > 2 else { return nil }

        let separatorRanges = regexMatches(Patterns.tableSeparator, in: line.source).map {
            NSRange(location: line.range.location + $0.range.location, length: $0.range.length)
        }
        guard separatorRanges.count >= 1 else { return nil }

        let cellRanges = tableCellRanges(in: line.source).map {
            NSRange(location: line.range.location + $0.location, length: $0.length)
        }
        return .tableRow(
            cellRanges: cellRanges,
            separatorRanges: separatorRanges,
            isDivider: isTableDivider(trimmed)
        )
    }

    private func inlineSpans(in line: EditorLine, kind: HybridMarkdownLineKind) -> [HybridMarkdownSpan] {
        guard line.range.length > 0 else { return [] }
        if case .horizontalRule = kind { return [] }
        if case .fencedCodeFence = kind { return [] }
        if case .codeBlockContent = kind { return [] }

        var spans: [HybridMarkdownSpan] = []
        var excludedRanges: [NSRange] = []

        if line.source.contains("`") {
            collectMatches(Patterns.inlineCode, in: line, excluding: excludedRanges).forEach { match in
                spans.append(HybridMarkdownSpan(range: match.contentRange, kind: .inlineCode))
                spans.append(contentsOf: delimiterRanges(match.fullRange, leading: 1, trailing: 1))
                excludedRanges.append(match.fullRange)
            }
        }

        if line.source.contains("**") {
            collectMatches(Patterns.boldAsterisk, in: line, excluding: excludedRanges).forEach { match in
                spans.append(HybridMarkdownSpan(range: match.contentRange, kind: .bold))
                spans.append(contentsOf: delimiterRanges(match.fullRange, leading: 2, trailing: 2))
                excludedRanges.append(match.fullRange)
            }
        }

        if line.source.contains("__") {
            collectMatches(Patterns.boldUnderscore, in: line, excluding: excludedRanges).forEach { match in
                spans.append(HybridMarkdownSpan(range: match.contentRange, kind: .bold))
                spans.append(contentsOf: delimiterRanges(match.fullRange, leading: 2, trailing: 2))
                excludedRanges.append(match.fullRange)
            }
        }

        if line.source.contains("](") {
            collectLinkMatches(in: line, excluding: excludedRanges).forEach { match in
                spans.append(HybridMarkdownSpan(range: match.titleRange, kind: .link))
                spans.append(HybridMarkdownSpan(range: match.urlRange, kind: .markdownDelimiter))
                spans.append(contentsOf: match.delimiterRanges.map {
                    HybridMarkdownSpan(range: $0, kind: .markdownDelimiter)
                })
                excludedRanges.append(match.fullRange)
            }
        }

        if !isReferenceStyleLinkDefinition(line.source),
           line.source.contains("https://") || line.source.contains("http://") {
            collectMatches(Patterns.automaticLink, in: line, excluding: excludedRanges).forEach { match in
                spans.append(HybridMarkdownSpan(range: match.fullRange, kind: .automaticLink))
                excludedRanges.append(match.fullRange)
            }
        }

        if line.source.contains("*") {
            collectMatches(Patterns.italicAsterisk, in: line, excluding: excludedRanges).forEach { match in
                spans.append(HybridMarkdownSpan(range: match.contentRange, kind: .italic))
                spans.append(contentsOf: delimiterRanges(match.fullRange, leading: 1, trailing: 1))
                excludedRanges.append(match.fullRange)
            }
        }

        if line.source.contains("_") {
            collectMatches(Patterns.italicUnderscore, in: line, excluding: excludedRanges).forEach { match in
                spans.append(HybridMarkdownSpan(range: match.contentRange, kind: .italic))
                spans.append(contentsOf: delimiterRanges(match.fullRange, leading: 1, trailing: 1))
                excludedRanges.append(match.fullRange)
            }
        }

        return spans.sorted {
            if $0.range.location == $1.range.location {
                return $0.range.length < $1.range.length
            }
            return $0.range.location < $1.range.location
        }
    }

    private func collectMatches(
        _ regex: NSRegularExpression,
        in line: EditorLine,
        excluding excludedRanges: [NSRange]
    ) -> [InlineMatch] {
        return regex.matches(in: line.source, range: NSRange(location: 0, length: line.source.utf16.count))
            .compactMap { match in
                guard match.numberOfRanges > 1 else { return nil }
                let shiftedRange = NSRange(
                    location: line.range.location + match.range.location,
                    length: match.range.length
                )
                guard !excludedRanges.contains(where: { $0.intersects(shiftedRange) }) else { return nil }
                return InlineMatch(
                    fullRange: shiftedRange,
                    contentRange: NSRange(
                        location: line.range.location + match.range(at: 1).location,
                        length: match.range(at: 1).length
                    )
                )
            }
    }

    private func collectLinkMatches(in line: EditorLine, excluding excludedRanges: [NSRange]) -> [LinkMatch] {
        Patterns.markdownLink.matches(in: line.source, range: NSRange(location: 0, length: line.source.utf16.count))
            .compactMap { match in
                guard match.numberOfRanges > 2 else { return nil }
                let fullRange = NSRange(location: line.range.location + match.range.location, length: match.range.length)
                guard !excludedRanges.contains(where: { $0.intersects(fullRange) }) else { return nil }
                let titleRange = NSRange(
                    location: line.range.location + match.range(at: 1).location,
                    length: match.range(at: 1).length
                )
                let urlRange = NSRange(
                    location: line.range.location + match.range(at: 2).location,
                    length: match.range(at: 2).length
                )
                return LinkMatch(
                    fullRange: fullRange,
                    titleRange: titleRange,
                    urlRange: urlRange,
                    delimiterRanges: [
                        NSRange(location: fullRange.location, length: 1),
                        NSRange(location: titleRange.upperBound, length: 2),
                        NSRange(location: fullRange.upperBound - 1, length: 1)
                    ]
                )
            }
    }

    private func delimiterRanges(
        _ fullRange: NSRange,
        leading: Int,
        trailing: Int
    ) -> [HybridMarkdownSpan] {
        [
            HybridMarkdownSpan(range: NSRange(location: fullRange.location, length: leading), kind: .markdownDelimiter),
            HybridMarkdownSpan(
                range: NSRange(location: fullRange.upperBound - trailing, length: trailing),
                kind: .markdownDelimiter
            )
        ]
    }

    private func firstMatch(_ regex: NSRegularExpression, in source: String) -> NSTextCheckingResult? {
        regexMatches(regex, in: source).first
    }

    private func regexMatches(_ regex: NSRegularExpression, in source: String) -> [NSTextCheckingResult] {
        return regex.matches(in: source, range: NSRange(location: 0, length: source.utf16.count))
    }

    private func nestingLevel(_ indentationLength: Int) -> Int {
        max(0, indentationLength / 2)
    }

    private func isHorizontalRule(_ source: String) -> Bool {
        let trimmed = source.trimmingCharacters(in: .whitespaces)
        guard trimmed.count >= 3 else { return false }
        let compact = trimmed.filter { $0 != " " && $0 != "\t" }
        guard compact.count >= 3, let first = compact.first, ["-", "*", "_"].contains(first) else {
            return false
        }
        return compact.allSatisfy { $0 == first }
    }

    private func isTableDivider(_ source: String) -> Bool {
        let cells = source.split(separator: "|", omittingEmptySubsequences: false)
            .map { $0.trimmingCharacters(in: .whitespaces) }
            .filter { !$0.isEmpty }
        guard !cells.isEmpty else { return false }
        return cells.allSatisfy { cell in
            let core = cell.trimmingCharacters(in: CharacterSet(charactersIn: ":"))
            return core.count >= 3 && core.allSatisfy { $0 == "-" }
        }
    }

    private func tableCellRanges(in source: String) -> [NSRange] {
        let nsSource = source as NSString
        var ranges: [NSRange] = []
        var cellStart = 0

        for location in 0..<nsSource.length where nsSource.character(at: location) == 124 {
            appendTrimmedRange(from: cellStart, to: location, in: nsSource, ranges: &ranges)
            cellStart = location + 1
        }
        appendTrimmedRange(from: cellStart, to: nsSource.length, in: nsSource, ranges: &ranges)

        return ranges.filter { $0.length > 0 }
    }

    private func appendTrimmedRange(from start: Int, to end: Int, in source: NSString, ranges: inout [NSRange]) {
        var lower = start
        var upper = end
        while lower < upper, isWhitespace(source.character(at: lower)) {
            lower += 1
        }
        while upper > lower, isWhitespace(source.character(at: upper - 1)) {
            upper -= 1
        }
        ranges.append(NSRange(location: lower, length: upper - lower))
    }

    private func isWhitespace(_ character: unichar) -> Bool {
        character == 32 || character == 9
    }

    private func isReferenceStyleLinkDefinition(_ source: String) -> Bool {
        guard firstNonWhitespaceCharacter(in: source) == "[",
              let closingBracket = source.firstIndex(of: "]")
        else { return false }

        let nextIndex = source.index(after: closingBracket)
        return nextIndex < source.endIndex && source[nextIndex] == ":"
    }

    private func firstNonWhitespaceCharacter(in source: String) -> Character? {
        source.first { $0 != " " && $0 != "\t" }
    }
}

private struct InlineMatch {
    var fullRange: NSRange
    var contentRange: NSRange
}

private struct LinkMatch {
    var fullRange: NSRange
    var titleRange: NSRange
    var urlRange: NSRange
    var delimiterRanges: [NSRange]
}

private enum Patterns {
    static let unorderedList = regex("^([ \\t]*)([-*])\\s+(.*)$")
    static let orderedList = regex("^([ \\t]*)(\\d+[.)])\\s+(.*)$")
    static let taskList = regex("^([ \\t]*)([-*])\\s+(\\[[ xX]\\])\\s+(.*)$")
    static let fencedCodeFence = regex("^([ \\t]*)(```|~~~)([^`~]*)$")
    static let tableSeparator = regex("\\|")
    static let inlineCode = regex("`([^`\\n]+)`")
    static let boldAsterisk = regex("\\*\\*([^*\\n]+)\\*\\*")
    static let boldUnderscore = regex("__([^_\\n]+)__")
    static let markdownLink = regex("(?<!!)\\[([^\\]\\n]+)\\]\\(([^)\\s\\n]+)\\)")
    static let automaticLink = regex("(https?://[^\\s<>()]+)")
    static let italicAsterisk = regex("(?<!\\*)\\*([^*\\n]+)\\*(?!\\*)")
    static let italicUnderscore = regex("(?<!_)_([^_\\n]+)_(?!_)")

    private static func regex(_ pattern: String) -> NSRegularExpression {
        do {
            return try NSRegularExpression(pattern: pattern)
        } catch {
            preconditionFailure("Invalid Markdown renderer regex: \(pattern)")
        }
    }
}

private extension NSRange {
    var upperBound: Int {
        location + length
    }

    func intersects(_ other: NSRange) -> Bool {
        location < other.upperBound && other.location < upperBound
    }
}