import AppKit

enum NotesVimMode: String {
    case normal = "NORMAL", insert = "INSERT", visual = "VISUAL", visualLine = "VISUAL LINE"
}

struct NotesVimState {
    var mode: NotesVimMode = .normal
    var count = 0
    var operatorCount = 1
    var pending: Character?
    var awaitingG = false
    var anchor: Int?
    var head: Int?
    var column: Int?
    var register: NSAttributedString?
    var linewiseRegister = false
}

// ponytail: scan the whole note for each command; cache character offsets if large notes make motions slow.
@MainActor
private struct VimBuffer {
    var characters: [Character] = []
    var ranges: [NSRange] = []
    let length: Int

    init(_ contents: NSAttributedString) {
        length = contents.length
        var offset = 0
        for character in contents.string {
            let size = String(character).utf16.count
            characters.append(character)
            ranges.append(NSRange(location: offset, length: size))
            offset += size
        }
    }

    var count: Int { characters.count }
    func offset(_ index: Int) -> Int { index < count ? ranges[max(0, index)].location : length }
    func index(_ offset: Int) -> Int { ranges.firstIndex { NSMaxRange($0) > offset } ?? count }
    func newline(_ index: Int) -> Bool { index < count && characters[index].isNewline }
    func lineStart(_ index: Int) -> Int {
        var result = min(index, count)
        while result > 0 && !newline(result - 1) { result -= 1 }
        return result
    }
    func lineEnd(_ index: Int) -> Int {
        var result = min(index, count)
        while result < count && !newline(result) { result += 1 }
        return result
    }
    func normal(_ index: Int) -> Int {
        let index = min(max(0, index), count)
        let start = lineStart(index), end = lineEnd(index)
        return end > start ? min(index, end - 1) : start
    }
    func firstNonblank(_ index: Int) -> Int {
        var result = lineStart(index)
        let end = lineEnd(index)
        while result < end && characters[result].isWhitespace { result += 1 }
        return normal(result)
    }
    func wordKind(_ index: Int, big: Bool) -> Int {
        guard index < count else { return 0 }
        let character = characters[index]
        if character.isWhitespace { return 0 }
        if big { return 1 }
        return character.isLetter || character.isNumber || character == "_" ? 1 : 2
    }
    func word(_ origin: Int, key: Character) -> Int {
        let big = key.isUppercase
        var index = min(origin, count)
        if key == "b" || key == "B" {
            if index > 0 { index -= 1 }
            while index > 0 && wordKind(index, big: big) == 0 { index -= 1 }
            let kind = wordKind(index, big: big)
            while index > 0 && wordKind(index - 1, big: big) == kind { index -= 1 }
        } else if key == "w" || key == "W" {
            let kind = wordKind(index, big: big)
            while index < count && wordKind(index, big: big) == kind { index += 1 }
            while index < count && wordKind(index, big: big) == 0 { index += 1 }
        } else {
            if index < count { index += 1 }
            while index < count && wordKind(index, big: big) == 0 { index += 1 }
            let kind = wordKind(index, big: big)
            while index + 1 < count && wordKind(index + 1, big: big) == kind { index += 1 }
            index = min(index, max(0, count - 1))
        }
        return index
    }
}

extension NotesTextView {
    var vimMode: NotesVimMode { vimState.mode }

    func resetVim() {
        endInsertionSession()
        let register = vimState.register, linewise = vimState.linewiseRegister
        vimState = NotesVimState()
        vimState.register = register
        vimState.linewiseRegister = linewise
        needsDisplay = true
        vimModeChanged?()
    }

    func clearVimCommand() {
        vimState.count = 0
        vimState.operatorCount = 1
        vimState.pending = nil
        vimState.awaitingG = false
    }

    private func setVimMode(_ mode: NotesVimMode) {
        if vimState.mode == .insert { endInsertionSession() }
        breakUndoCoalescing()
        clearVimCommand()
        vimState.mode = mode
        vimState.anchor = nil
        vimState.head = nil
        vimState.column = nil
        if mode == .insert { beginInsertionSession() }
        needsDisplay = true
        vimModeChanged?()
    }

    @discardableResult
    func handleVimEscape() -> Bool {
        guard vimEnabled, vimState.mode != .normal, !hasMarkedText() else { return false }
        let buffer = VimBuffer(documentContents)
        var head = vimState.head ?? buffer.index(selectedRange().location)
        if vimState.mode == .insert && head > buffer.lineStart(head) { head -= 1 }
        setVimMode(.normal)
        setVimCaret(buffer.normal(head), buffer: buffer)
        return true
    }

    private func setVimCaret(_ index: Int, buffer: VimBuffer) {
        setSelectedRange(NSRange(location: buffer.offset(index), length: 0), affinity: .downstream, stillSelecting: false)
        scrollRangeToVisible(selectedRange())
    }

    private func showVimSelection(_ head: Int, buffer: VimBuffer) {
        let anchor = vimState.anchor ?? head
        vimState.head = head
        let low = min(anchor, head), high = max(anchor, head)
        let range: NSRange
        if vimState.mode == .visualLine {
            let start = buffer.offset(buffer.lineStart(low))
            let end = buffer.lineEnd(high)
            let stop = buffer.offset(min(buffer.count, end + (buffer.newline(end) ? 1 : 0)))
            range = NSRange(location: start, length: stop - start)
        } else {
            let start = buffer.offset(low)
            let stop = high < buffer.count ? NSMaxRange(buffer.ranges[high]) : buffer.length
            range = NSRange(location: start, length: stop - start)
        }
        setSelectedRange(range)
        scrollRangeToVisible(NSRange(location: buffer.offset(head), length: 0))
    }

    func handleVimKey(_ event: NSEvent) -> Bool {
        guard vimEnabled, isEditable, !hasMarkedText() else { return false }
        let flags = event.modifierFlags.intersection([.command, .control, .option, .shift])
        // Command shortcuts remain native; Option text input belongs only to Insert mode.
        if flags.contains(.command) { clearVimCommand(); return false }
        if flags.contains(.option) { clearVimCommand(); return vimState.mode != .insert }
        if event.keyCode == 53 { return handleVimEscape() }
        if flags == .control {
            if event.charactersIgnoringModifiers == "[" { return handleVimEscape() }
            if vimState.mode != .insert && event.charactersIgnoringModifiers?.lowercased() == "r" {
                clearVimCommand()
                undoManager?.redo()
                normalizeVimCaret()
                return true
            }
            return vimState.mode != .insert
        }
        guard vimState.mode != .insert else { return false }
        let key: Character
        switch event.keyCode {
        case 123: key = "h"
        case 124: key = "l"
        case 125: key = "j"
        case 126: key = "k"
        case 51: key = "h"
        case 36, 76: key = "j"
        default:
            guard let text = event.characters, text.count == 1, let character = text.first else { return true }
            key = character
        }
        if let digit = key.wholeNumberValue, digit < 10, digit > 0 || vimState.count > 0 {
            // Saturate before arithmetic; a held numeric key must not overflow or create unbounded work.
            vimState.count = min(9999, vimState.count * 10 + digit)
            return true
        }
        let buffer = VimBuffer(documentContents)
        if vimState.mode == .normal, selectedRange().length > 0 {
            vimState.mode = .visual
            vimState.anchor = buffer.index(selectedRange().location)
            vimState.head = max(vimState.anchor!, buffer.index(NSMaxRange(selectedRange())) - 1)
            needsDisplay = true
            vimModeChanged?()
        }
        let origin = buffer.normal(vimState.head ?? buffer.index(selectedRange().location))
        let count = min(9999, max(1, vimState.count) * vimState.operatorCount)
        let visual = vimState.mode == .visual || vimState.mode == .visualLine

        if vimState.awaitingG {
            vimState.awaitingG = false
            guard key == "g" else { clearVimCommand(); return true }
            let target = vimLine(count: count, buffer: buffer)
            finishVimMotion(origin: origin, target: target, inclusive: false, linewise: true, buffer: buffer)
            return true
        }
        if key == "g" { vimState.awaitingG = true; return true }
        if "dcy".contains(key) {
            if visual {
                let range = selectedRange()
                applyVimOperator(key, range: range, linewise: vimState.mode == .visualLine)
            } else if vimState.pending == key {
                var end = buffer.lineEnd(origin)
                for _ in 1..<count where buffer.newline(end) { end = buffer.lineEnd(end + 1) }
                applyVimLineOperator(key, start: origin, end: end, buffer: buffer)
            } else if vimState.pending == nil {
                vimState.pending = key
                vimState.operatorCount = max(1, vimState.count)
                vimState.count = 0
            } else { clearVimCommand() }
            return true
        }
        if "hljkwWbBeE0^$G".contains(key) {
            var target = origin
            var inclusive = false
            let linewise = key == "j" || key == "k" || key == "G"
            switch key {
            case "h": target = max(buffer.lineStart(origin), origin - count)
            case "l": target = min(buffer.lineEnd(origin) - (vimState.pending == nil ? 1 : 0), origin + count)
                target = max(buffer.lineStart(origin), target)
            case "j", "k":
                let column = vimState.column ?? (origin - buffer.lineStart(origin))
                vimState.column = column
                var start = buffer.lineStart(origin)
                for _ in 0..<count {
                    if key == "j" {
                        let end = buffer.lineEnd(start)
                        if buffer.newline(end) { start = end + 1 }
                    } else if start > 0 { start = buffer.lineStart(start - 1) }
                }
                target = buffer.normal(min(start + column, buffer.lineEnd(start)))
            case "0": target = buffer.lineStart(origin)
            case "^": target = buffer.firstNonblank(origin)
            case "$":
                var end = buffer.lineEnd(origin)
                for _ in 1..<count where buffer.newline(end) { end = buffer.lineEnd(end + 1) }
                target = max(buffer.lineStart(end), end - 1)
                inclusive = true
            case "G": target = vimState.count > 0 || vimState.operatorCount > 1
                ? vimLine(count: count, buffer: buffer) : buffer.firstNonblank(buffer.count)
            default:
                var motion = key
                var repetitions = count
                // cw includes this word's end even when the cursor is already on its last character.
                if vimState.pending == "c", (key == "w" || key == "W"), buffer.wordKind(origin, big: key == "W") != 0 {
                    motion = key == "w" ? "e" : "E"
                    let kind = buffer.wordKind(origin, big: key == "W")
                    while target + 1 < buffer.count && buffer.wordKind(target + 1, big: key == "W") == kind { target += 1 }
                    repetitions -= 1
                }
                for step in 0..<repetitions {
                    let previous = target
                    target = buffer.word(target, key: motion)
                    if step == repetitions - 1, vimState.pending != nil, motion == "w" || motion == "W",
                       buffer.wordKind(previous, big: motion == "W") != 0, target > buffer.lineEnd(previous) {
                        target = buffer.lineEnd(previous)
                    }
                }
                inclusive = motion == "e" || motion == "E"
            }
            if key != "j" && key != "k" { vimState.column = nil }
            finishVimMotion(origin: origin, target: target, inclusive: inclusive, linewise: linewise, buffer: buffer)
            return true
        }
        guard vimState.pending == nil else { clearVimCommand(); return true }
        clearVimCommand()
        switch key {
        case "v", "V":
            let mode: NotesVimMode = key == "v" ? .visual : .visualLine
            if vimState.mode == mode { setVimMode(.normal); setVimCaret(origin, buffer: buffer) }
            else {
                let anchor = vimState.anchor ?? origin
                setVimMode(mode)
                vimState.anchor = anchor
                showVimSelection(origin, buffer: buffer)
            }
        case "i", "a", "I", "A":
            guard !visual else { return true }
            var target = origin
            if key == "a" && origin < buffer.lineEnd(origin) { target += 1 }
            if key == "I" { target = buffer.firstNonblank(origin) }
            if key == "A" { target = buffer.lineEnd(origin) }
            setVimMode(.insert)
            setVimCaret(target, buffer: buffer)
        case "o", "O":
            guard !visual else { return true }
            let target = key == "o" ? buffer.lineEnd(origin) : buffer.lineStart(origin)
            setVimMode(.insert)
            setVimCaret(target, buffer: buffer)
            insertNewline(nil)
            if key == "O" {
                let updated = VimBuffer(documentContents)
                let caret = updated.index(selectedRange().location)
                setVimCaret(updated.lineStart(max(0, caret - 1)), buffer: updated)
            }
        case "x":
            if visual { applyVimOperator("d", range: selectedRange(), linewise: vimState.mode == .visualLine) }
            else {
                let end = min(buffer.lineEnd(origin), origin + count)
                let startOffset = buffer.offset(origin), endOffset = buffer.offset(end)
                if end > origin { applyVimOperator("d", range: NSRange(location: startOffset, length: endOffset - startOffset), linewise: false) }
            }
        case "D", "C":
            var end = buffer.lineEnd(origin)
            for _ in 1..<count where buffer.newline(end) { end = buffer.lineEnd(end + 1) }
            let startOffset = buffer.offset(origin), endOffset = buffer.offset(end)
            applyVimOperator(key == "D" ? "d" : "c", range: NSRange(location: startOffset, length: endOffset - startOffset), linewise: false)
        case "p", "P": pasteVimRegister(before: key == "P", count: count, origin: origin, buffer: buffer)
        case "u":
            setVimMode(.normal)
            for _ in 0..<count where undoManager?.canUndo == true { undoManager?.undo() }
            normalizeVimCaret()
        default: break // Unsupported Normal-mode keys never become note text.
        }
        return true
    }

    private func vimLine(count: Int, buffer: VimBuffer) -> Int {
        var start = 0
        for _ in 1..<count {
            let end = buffer.lineEnd(start)
            if buffer.newline(end) { start = end + 1 }
        }
        return buffer.firstNonblank(start)
    }

    private func finishVimMotion(origin: Int, target: Int, inclusive: Bool, linewise: Bool, buffer: VimBuffer) {
        if let operation = vimState.pending {
            if linewise { applyVimLineOperator(operation, start: min(origin, target), end: max(origin, target), buffer: buffer) }
            else {
                let low = min(origin, target)
                let high = min(buffer.count, max(origin, target) + (inclusive ? 1 : 0))
                let start = buffer.offset(low), end = buffer.offset(high)
                // A motion against a boundary is a no-op, not a character deletion.
                if end > start { applyVimOperator(operation, range: NSRange(location: start, length: end - start), linewise: false) }
                else { clearVimCommand() }
            }
        } else {
            clearVimCommand()
            if vimState.mode == .visual || vimState.mode == .visualLine { showVimSelection(buffer.normal(target), buffer: buffer) }
            else { setVimCaret(buffer.normal(target), buffer: buffer) }
        }
    }

    private func applyVimLineOperator(_ operation: Character, start: Int, end: Int, buffer: VimBuffer) {
        let first = buffer.offset(buffer.lineStart(start))
        let lastLineEnd = buffer.lineEnd(end)
        let last = buffer.offset(lastLineEnd + (buffer.newline(lastLineEnd) ? 1 : 0))
        applyVimOperator(operation, range: NSRange(location: first, length: last - first), linewise: true)
    }

    private func applyVimOperator(_ operation: Character, range: NSRange, linewise: Bool) {
        let source = documentContents
        let register = NSMutableAttributedString(attributedString: source.attributedSubstring(from: range))
        if linewise && !register.string.hasSuffix("\n") {
            var attributes = register.length > 0 ? register.attributes(at: register.length - 1, effectiveRange: nil) : [:]
            attributes[.attachment] = nil
            register.append(NSAttributedString(string: "\n", attributes: attributes))
        }
        vimState.register = register
        vimState.linewiseRegister = linewise
        let originalCaret = selectedRange().location
        setVimMode(operation == "c" ? .insert : .normal)
        if operation == "y" {
            setSelectedRange(NSRange(location: originalCaret, length: 0))
            normalizeVimCaret()
            return
        }
        var deletion = range
        var replacement = NSAttributedString(string: "")
        var attributes = range.location < source.length ? source.attributes(at: range.location, effectiveRange: nil) : typingAttributes
        attributes[.attachment] = nil
        if linewise {
            let text = source.string as NSString
            if operation == "c" {
                // cc keeps the following line separate while opening a blank line for insertion.
                if NSMaxRange(range) < text.length { replacement = NSAttributedString(string: "\n", attributes: attributes) }
            } else if NSMaxRange(range) == text.length, range.location > 0 {
                let previous = text.rangeOfComposedCharacterSequence(at: range.location - 1)
                if text.substring(with: previous).allSatisfy(\.isNewline) {
                    deletion.location = previous.location
                    deletion.length += previous.length
                }
            }
        }
        replaceDocument(replacement, range: deletion, selection: NSRange(location: deletion.location, length: 0))
        if operation == "c" { typingAttributes = attributes }
        else { normalizeVimCaret() }
    }

    private func pasteVimRegister(before: Bool, count: Int, origin: Int, buffer: VimBuffer) {
        guard let register = vimState.register else { return }
        let insertion = NSMutableAttributedString(string: "")
        for _ in 0..<count { insertion.append(register) }
        var range: NSRange
        if vimState.mode == .visual || vimState.mode == .visualLine {
            range = selectedRange()
        } else if vimState.linewiseRegister {
            let source = documentContents.string as NSString
            let line = source.lineRange(for: NSRange(location: buffer.offset(origin), length: 0))
            let offset = before ? line.location : NSMaxRange(line)
            range = NSRange(location: offset, length: 0)
            if offset == source.length, offset == 0 || !source.substring(from: offset - 1).allSatisfy(\.isNewline) {
                if offset > 0 {
                    insertion.insert(NSAttributedString(string: "\n"), at: 0)
                }
                if insertion.string.hasSuffix("\n") { insertion.deleteCharacters(in: NSRange(location: insertion.length - 1, length: 1)) }
            }
        } else {
            let index = before ? origin : min(buffer.lineEnd(origin), origin + 1)
            range = NSRange(location: buffer.offset(index), length: 0)
        }
        setVimMode(.normal)
        replaceDocument(insertion, range: range)
        let caret = vimState.linewiseRegister ? range.location + (insertion.string.hasPrefix("\n") ? 1 : 0) : range.location + insertion.length
        setSelectedRange(NSRange(location: caret, length: 0))
        if !vimState.linewiseRegister {
            let updated = VimBuffer(documentContents)
            setVimCaret(updated.normal(max(0, updated.index(selectedRange().location) - 1)), buffer: updated)
        } else { normalizeVimCaret() }
    }

    private func normalizeVimCaret() {
        let buffer = VimBuffer(documentContents)
        setVimCaret(buffer.normal(buffer.index(selectedRange().location)), buffer: buffer)
        vimState.anchor = nil
        vimState.head = nil
    }
}