Sapling/Sources/SaplingEditor/DocumentPresentationState.swift

import Foundation

public enum RenderedLineState: String, Hashable, Sendable {
    case source
    case rendered
}

public struct RenderedHeadingElement: Hashable, Sendable {
    public var lineIndex: Int
    public var level: Int
    public var markerRange: NSRange
    public var textRange: NSRange
}

public struct RenderedTaskElement: Hashable, Sendable {
    public var lineIndex: Int
    public var markerRange: NSRange
    public var checkboxRange: NSRange
    public var contentRange: NSRange
    public var checked: Bool
    public var nestingLevel: Int

    public var toggledMarkdownCheckbox: String {
        checked ? "[ ]" : "[x]"
    }
}

public struct RenderedCodeBlockElement: Hashable, Sendable {
    public var lineIndexes: [Int]
    public var sourceRange: NSRange
    public var languageRange: NSRange?
}

public struct RenderedLinkElement: Hashable, Sendable {
    public var lineIndex: Int
    public var titleRange: NSRange
    public var urlRange: NSRange?
}

public enum RenderedDocumentElement: Hashable, Sendable {
    case heading(RenderedHeadingElement)
    case task(RenderedTaskElement)
    case codeBlock(RenderedCodeBlockElement)
    case link(RenderedLinkElement)
    case inlineCode(lineIndex: Int, range: NSRange)
    case blockquote(lineIndex: Int, range: NSRange)
    case horizontalRule(lineIndex: Int, range: NSRange)
    case listItem(lineIndex: Int, markerRange: NSRange, contentRange: NSRange, nestingLevel: Int)
    case tableRow(lineIndex: Int, range: NSRange)
    case paragraph(lineIndex: Int, range: NSRange)
}

public struct RenderNode: Hashable, Sendable {
    public var element: RenderedDocumentElement

    public init(element: RenderedDocumentElement) {
        self.element = element
    }
}

public struct RenderedLine: Hashable, Sendable {
    public var line: EditorLine
    public var renderPlan: HybridMarkdownLineRenderPlan
    public var nodes: [RenderNode]
}

public struct RenderSnapshot: Hashable, Sendable {
    public var lines: [String]

    public init(lines: [String]) {
        self.lines = lines
    }

    public var signature: String {
        lines.joined(separator: "\n")
    }
}

public struct DocumentRenderModel: Hashable, Sendable {
    public var lineCount: Int
    public var lines: [RenderedLine]
    public var nodes: [RenderNode]

    public init(lineIndex: DocumentLineIndex, lineIndexes: [Int]? = nil) {
        self.lineCount = lineIndex.lineCount

        let renderer = HybridMarkdownLineRenderer()
        let selectedLines = lineIndexes.map {
            lineIndex.editorLines(for: $0, activeLineIndex: -1)
        } ?? lineIndex.editorLines(activeLineIndex: -1)
        let renderPlans = DocumentPresentationState.renderPlans(
            for: selectedLines,
            lineIndex: lineIndex,
            activeLineIndex: -1,
            renderer: renderer,
            needsDocumentContext: lineIndexes != nil
        )
        let renderPlansByLine = Dictionary(uniqueKeysWithValues: renderPlans.map { ($0.line.index, $0) })
        var renderedLines: [RenderedLine] = []
        var collectedNodes: [RenderNode] = []

        for line in selectedLines {
            let renderPlan = renderPlansByLine[line.index] ?? renderer.renderPlan(for: line)
            let nodes = DocumentPresentationState.elements(for: renderPlan, state: .rendered).map(RenderNode.init)
            renderedLines.append(RenderedLine(line: line, renderPlan: renderPlan, nodes: nodes))
            collectedNodes.append(contentsOf: nodes)
        }

        if lineIndexes == nil {
            collectedNodes.append(contentsOf: DocumentPresentationState.codeBlockElements(from: renderPlans).map(RenderNode.init))
        }

        self.lines = renderedLines
        self.nodes = collectedNodes
    }

    public var snapshot: RenderSnapshot {
        let lineSnapshots = lines.map { renderedLine in
            [
                "line=\(renderedLine.line.index)",
                "range=\(renderedLine.line.range.location):\(renderedLine.line.range.length)",
                "kind=\(Self.describe(renderedLine.renderPlan.kind))",
                "spans=\(renderedLine.renderPlan.spans.map(Self.describe).joined(separator: ","))",
                "nodes=\(renderedLine.nodes.map { Self.describe($0.element) }.joined(separator: ","))"
            ].joined(separator: "|")
        }
        let treeSnapshot = "tree=" + nodes.map { Self.describe($0.element) }.joined(separator: "|")
        return RenderSnapshot(lines: lineSnapshots + [treeSnapshot])
    }

    private static func describe(_ kind: HybridMarkdownLineKind) -> String {
        switch kind {
        case .paragraph:
            return "paragraph"
        case .heading(let level, let markerRange, let textRange):
            return "heading:\(level):\(describe(markerRange)):\(describe(textRange))"
        case .blockquote(let markerRange, let contentRange):
            return "blockquote:\(describe(markerRange)):\(describe(contentRange))"
        case .horizontalRule(let range):
            return "horizontalRule:\(describe(range))"
        case .unorderedList(let markerRange, let contentRange, let nestingLevel):
            return "unorderedList:\(describe(markerRange)):\(describe(contentRange)):\(nestingLevel)"
        case .orderedList(let markerRange, let contentRange, let nestingLevel):
            return "orderedList:\(describe(markerRange)):\(describe(contentRange)):\(nestingLevel)"
        case .taskList(let markerRange, let checkboxRange, let contentRange, let checked, let nestingLevel):
            return "task:\(describe(markerRange)):\(describe(checkboxRange)):\(describe(contentRange)):\(checked):\(nestingLevel)"
        case .fencedCodeFence(let markerRange, let languageRange):
            return "codeFence:\(describe(markerRange)):\(languageRange.map(describe) ?? "nil")"
        case .codeBlockContent(let language):
            return "codeContent:\(language ?? "plain")"
        case .tableRow(let cellRanges, let separatorRanges, let isDivider):
            return "table:\(cellRanges.map(describe).joined(separator: ",")):"
                + "\(separatorRanges.map(describe).joined(separator: ",")):\(isDivider)"
        }
    }

    private static func describe(_ span: HybridMarkdownSpan) -> String {
        "\(span.kind):\(describe(span.range))"
    }

    private static func describe(_ element: RenderedDocumentElement) -> String {
        switch element {
        case .heading(let heading):
            return "heading:\(heading.lineIndex):\(heading.level):\(describe(heading.textRange))"
        case .task(let task):
            return "task:\(task.lineIndex):\(describe(task.checkboxRange)):\(task.checked)"
        case .codeBlock(let codeBlock):
            return "codeBlock:\(codeBlock.lineIndexes.map(String.init).joined(separator: ",")):"
                + "\(describe(codeBlock.sourceRange)):\(codeBlock.languageRange.map(describe) ?? "nil")"
        case .link(let link):
            return "link:\(link.lineIndex):\(describe(link.titleRange)):\(link.urlRange.map(describe) ?? "nil")"
        case .inlineCode(let lineIndex, let range):
            return "inlineCode:\(lineIndex):\(describe(range))"
        case .blockquote(let lineIndex, let range):
            return "blockquote:\(lineIndex):\(describe(range))"
        case .horizontalRule(let lineIndex, let range):
            return "horizontalRule:\(lineIndex):\(describe(range))"
        case .listItem(let lineIndex, let markerRange, let contentRange, let nestingLevel):
            return "listItem:\(lineIndex):\(describe(markerRange)):\(describe(contentRange)):\(nestingLevel)"
        case .tableRow(let lineIndex, let range):
            return "tableRow:\(lineIndex):\(describe(range))"
        case .paragraph(let lineIndex, let range):
            return "paragraph:\(lineIndex):\(describe(range))"
        }
    }

    private static func describe(_ range: NSRange) -> String {
        "\(range.location):\(range.length)"
    }
}

public struct DocumentPresentationLine: Hashable, Sendable {
    public var line: EditorLine
    public var state: RenderedLineState
    public var renderPlan: HybridMarkdownLineRenderPlan
    public var elements: [RenderedDocumentElement]
}

public struct DocumentPresentationState: Hashable, Sendable {
    public var activeLineIndex: Int
    public var editableRegion: EditableRegion
    public var lineCount: Int
    public var lines: [DocumentPresentationLine]
    public var elements: [RenderedDocumentElement]

    public init(
        lineIndex: DocumentLineIndex,
        activeLineIndex: Int,
        lineIndexes: [Int]? = nil
    ) {
        self.init(
            lineIndex: lineIndex,
            editableRegion: activeLineIndex >= 0 ? EditableRegion(lineIndexes: [activeLineIndex]) : .none(),
            lineIndexes: lineIndexes
        )
    }

    public init(
        lineIndex: DocumentLineIndex,
        editableRegion: EditableRegion,
        lineIndexes: [Int]? = nil
    ) {
        self.activeLineIndex = editableRegion.primaryLineIndex
        self.editableRegion = editableRegion
        self.lineCount = lineIndex.lineCount

        let renderer = HybridMarkdownLineRenderer()
        let selectedLines = lineIndexes.map {
            lineIndex.editorLines(for: $0, activeLineIndex: editableRegion.primaryLineIndex)
        } ?? lineIndex.editorLines(activeLineIndex: editableRegion.primaryLineIndex)
        let renderPlans = Self.renderPlans(
            for: selectedLines,
            lineIndex: lineIndex,
            activeLineIndex: editableRegion.primaryLineIndex,
            renderer: renderer,
            needsDocumentContext: lineIndexes != nil
        )
        let renderPlansByLine = Dictionary(uniqueKeysWithValues: renderPlans.map { ($0.line.index, $0) })
        var presentationLines: [DocumentPresentationLine] = []
        var collectedElements: [RenderedDocumentElement] = []

        for line in selectedLines {
            let state: RenderedLineState = editableRegion.contains(line.index) ? .source : .rendered
            let renderPlan = renderPlansByLine[line.index] ?? renderer.renderPlan(for: line)
            let lineElements = Self.elements(for: renderPlan, state: state)
            presentationLines.append(DocumentPresentationLine(
                line: line,
                state: state,
                renderPlan: renderPlan,
                elements: lineElements
            ))
            collectedElements.append(contentsOf: lineElements)
        }

        if lineIndexes == nil {
            collectedElements.append(contentsOf: Self.codeBlockElements(from: renderPlans))
        }

        self.lines = presentationLines
        self.elements = collectedElements
    }

    public var renderedTasks: [RenderedTaskElement] {
        elements.compactMap {
            guard case .task(let task) = $0 else { return nil }
            return task
        }
    }

    public var renderedCodeBlocks: [RenderedCodeBlockElement] {
        elements.compactMap {
            guard case .codeBlock(let codeBlock) = $0 else { return nil }
            return codeBlock
        }
    }

    public func lineState(at lineIndex: Int) -> RenderedLineState? {
        lines.first { $0.line.index == lineIndex }?.state
    }

    public static func renderedTasks(
        in lineIndex: DocumentLineIndex,
        activeLineIndex: Int
    ) -> [RenderedTaskElement] {
        guard lineIndex.source.contains("[ ]")
            || lineIndex.source.contains("[x]")
            || lineIndex.source.contains("[X]")
        else {
            return []
        }

        return DocumentPresentationState(
            lineIndex: lineIndex,
            activeLineIndex: activeLineIndex
        ).renderedTasks
    }

    fileprivate static func renderPlans(
        for lines: [EditorLine],
        lineIndex: DocumentLineIndex,
        activeLineIndex: Int,
        renderer: HybridMarkdownLineRenderer,
        needsDocumentContext: Bool
    ) -> [HybridMarkdownLineRenderPlan] {
        if needsDocumentContext {
            guard linesNeedCodeBlockContext(lines, lineIndex: lineIndex) else {
                return renderer.renderPlans(for: lines)
            }
            return renderer.renderPlans(
                for: lines,
                resolvingCodeBlockContextWith: lineIndex,
                activeLineIndex: activeLineIndex
            )
        }
        return renderer.renderPlans(for: lines)
    }

    private static func linesNeedCodeBlockContext(_ lines: [EditorLine], lineIndex: DocumentLineIndex) -> Bool {
        for line in lines {
            let lowerBound = max(0, line.index - 2)
            let upperBound = min(lineIndex.lineCount - 1, line.index + 2)
            for nearbyLineIndex in lowerBound...upperBound {
                guard let nearbyLine = lineIndex.editorLine(at: nearbyLineIndex, activeLineIndex: -1) else { continue }
                if isFenceLine(nearbyLine.source) {
                    return true
                }
            }
        }
        return false
    }

    private static func isFenceLine(_ source: String) -> Bool {
        let trimmedPrefix = source.prefix { $0 == " " || $0 == "\t" }
        guard trimmedPrefix.count <= 3 else { return false }
        let content = source.dropFirst(trimmedPrefix.count)
        return content.hasPrefix("```") || content.hasPrefix("~~~")
    }

    fileprivate static func elements(
        for renderPlan: HybridMarkdownLineRenderPlan,
        state: RenderedLineState
    ) -> [RenderedDocumentElement] {
        guard state == .rendered else { return [] }

        var elements: [RenderedDocumentElement] = []
        switch renderPlan.kind {
        case .heading(let level, let markerRange, let textRange):
            elements.append(.heading(RenderedHeadingElement(
                lineIndex: renderPlan.line.index,
                level: level,
                markerRange: markerRange,
                textRange: textRange
            )))
        case .blockquote:
            elements.append(.blockquote(lineIndex: renderPlan.line.index, range: renderPlan.line.range))
        case .horizontalRule(let range):
            elements.append(.horizontalRule(lineIndex: renderPlan.line.index, range: range))
        case .unorderedList(let markerRange, let contentRange, let nestingLevel),
             .orderedList(let markerRange, let contentRange, let nestingLevel):
            elements.append(.listItem(
                lineIndex: renderPlan.line.index,
                markerRange: markerRange,
                contentRange: contentRange,
                nestingLevel: nestingLevel
            ))
        case .taskList(let markerRange, let checkboxRange, let contentRange, let checked, let nestingLevel):
            elements.append(.task(RenderedTaskElement(
                lineIndex: renderPlan.line.index,
                markerRange: markerRange,
                checkboxRange: checkboxRange,
                contentRange: contentRange,
                checked: checked,
                nestingLevel: nestingLevel
            )))
        case .fencedCodeFence, .codeBlockContent:
            break
        case .tableRow:
            elements.append(.tableRow(lineIndex: renderPlan.line.index, range: renderPlan.line.range))
        case .paragraph:
            break
        }

        for span in renderPlan.spans {
            switch span.kind {
            case .inlineCode:
                elements.append(.inlineCode(lineIndex: renderPlan.line.index, range: span.range))
            case .link:
                elements.append(.link(RenderedLinkElement(
                    lineIndex: renderPlan.line.index,
                    titleRange: span.range,
                    urlRange: nil
                )))
            case .automaticLink:
                elements.append(.link(RenderedLinkElement(
                    lineIndex: renderPlan.line.index,
                    titleRange: span.range,
                    urlRange: span.range
                )))
            case .bold, .italic, .markdownDelimiter:
                break
            }
        }

        return elements
    }

    fileprivate static func codeBlockElements(from plans: [HybridMarkdownLineRenderPlan]) -> [RenderedDocumentElement] {
        var elements: [RenderedDocumentElement] = []
        var openLineIndexes: [Int] = []
        var openRange: NSRange?
        var openLanguageRange: NSRange?

        for plan in plans {
            guard case .fencedCodeFence(let markerRange, let languageRange) = plan.kind else {
                if openRange != nil {
                    openLineIndexes.append(plan.line.index)
                }
                continue
            }

            if let sourceRange = openRange {
                openLineIndexes.append(plan.line.index)
                elements.append(.codeBlock(RenderedCodeBlockElement(
                    lineIndexes: openLineIndexes,
                    sourceRange: NSRange(
                        location: sourceRange.location,
                        length: plan.line.range.upperBound - sourceRange.location
                    ),
                    languageRange: openLanguageRange
                )))
                openLineIndexes.removeAll()
                openRange = nil
                openLanguageRange = nil
            } else {
                openLineIndexes = [plan.line.index]
                openRange = markerRange
                openLanguageRange = languageRange
            }
        }

        if let sourceRange = openRange,
           let lastLine = plans.last?.line {
            elements.append(.codeBlock(RenderedCodeBlockElement(
                lineIndexes: openLineIndexes,
                sourceRange: NSRange(
                    location: sourceRange.location,
                    length: lastLine.range.upperBound - sourceRange.location
                ),
                languageRange: openLanguageRange
            )))
        }

        return elements
    }
}

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