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 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 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 lineCount: Int
    public var lines: [DocumentPresentationLine]
    public var elements: [RenderedDocumentElement]

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

        let renderer = HybridMarkdownLineRenderer()
        let selectedLines = lineIndexes.map {
            lineIndex.editorLines(for: $0, activeLineIndex: activeLineIndex)
        } ?? lineIndex.editorLines(activeLineIndex: activeLineIndex)
        let renderPlans = Self.renderPlans(
            for: selectedLines,
            lineIndex: lineIndex,
            activeLineIndex: activeLineIndex,
            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 = line.index == activeLineIndex ? .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
    }

    private 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("~~~")
    }

    private 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
    }

    private 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
    }
}