Tusker/Packages/ComposeUI/Sources/ComposeUI/UITextInput+Autocomplete.swift

//
//  UITextInput+Autocomplete.swift
//  ComposeUI
//
//  Created by Shadowfacts on 3/5/23.
//

import UIKit
import SwiftUI

extension UITextInput {
    func autocomplete(with string: String, permittedModes: AutocompleteModes, autocompleteState: inout AutocompleteState?) {
        guard let selectedTextRange,
              let wholeDocumentRange = self.textRange(from: self.beginningOfDocument, to: self.endOfDocument),
              let text = self.text(in: wholeDocumentRange),
              let (lastWordStartIndex, _) = findAutocompleteLastWord() else {
            return
        }
        
        let distanceToEnd = self.offset(from: selectedTextRange.start, to: self.endOfDocument)
        
        let selectedRangeStartUTF16 = self.offset(from: self.beginningOfDocument, to: selectedTextRange.start)
        let characterBeforeCursorIndex = text.utf16.index(text.startIndex, offsetBy: selectedRangeStartUTF16)
        
        let insertSpace: Bool
        if distanceToEnd > 0 {
            let charAfterCursor = text[characterBeforeCursorIndex]
            insertSpace = charAfterCursor != " " && charAfterCursor != "\n"
        } else {
            insertSpace = true
        }
        let string = insertSpace ? string + " " : string
        
        let startPosition = self.position(from: self.beginningOfDocument, offset: text.utf16.distance(from: text.startIndex, to: lastWordStartIndex))!
        let lastWordRange = self.textRange(from: startPosition, to: selectedTextRange.start)!
        replace(lastWordRange, withText: string)
        
        autocompleteState = updateAutocompleteState(permittedModes: permittedModes)
        
        // keep the cursor at the same position in the text, immediately after what was inserted
        // if we inserted a space, move the cursor 1 farther so it's immediately after the pre-existing space
        let insertSpaceOffset = insertSpace ? 0 : 1
        let newCursorPosition = self.position(from: self.endOfDocument, offset: -distanceToEnd + insertSpaceOffset)!
        self.selectedTextRange = self.textRange(from: newCursorPosition, to: newCursorPosition)
    }
    
    func updateAutocompleteState(permittedModes: AutocompleteModes) -> AutocompleteState? {
        guard let selectedTextRange,
            let wholeDocumentRange = self.textRange(from: self.beginningOfDocument, to: self.endOfDocument),
              let text = self.text(in: wholeDocumentRange),
              !text.isEmpty,
              let (lastWordStartIndex, foundFirstAtSign) = findAutocompleteLastWord() else {
            return nil
        }
        
        let triggerChars = permittedModes.triggerChars
        
        if lastWordStartIndex > text.startIndex {
            // if the character before the "word" beginning is a valid part of a "word",
            // we aren't able to autocomplete
            let c = text[text.index(before: lastWordStartIndex)]
            if isPermittedForAutocomplete(c) || triggerChars.contains(c) {
                return nil
            }
        }
        
        let characterBeforeCursorIndex = text.utf16.index(text.startIndex, offsetBy: self.offset(from: self.beginningOfDocument, to: selectedTextRange.start))
        
        if lastWordStartIndex >= text.startIndex {
            let lastWord = text[lastWordStartIndex..<characterBeforeCursorIndex]
            let exceptFirst = lastWord[lastWord.index(after: lastWord.startIndex)...]
            
            // periods are only allowed in mentions in the domain part
            if lastWord.contains(".") {
                if lastWord.first == "@" && foundFirstAtSign && permittedModes.contains(.mentions) {
                    return .mention(String(exceptFirst))
                } else {
                    return nil
                }
            }
            
            switch lastWord.first {
            case "@" where permittedModes.contains(.mentions):
                return .mention(String(exceptFirst))
            case ":" where permittedModes.contains(.emojis):
                return .emoji(String(exceptFirst))
            case "#" where permittedModes.contains(.hashtags):
                return .hashtag(String(exceptFirst))
            default:
                return nil
            }
        } else {
            return nil
        }
    }
    
    private func findAutocompleteLastWord() -> (index: String.Index, foundFirstAtSign: Bool)? {
        guard (self as? UIView)?.isFirstResponder == true,
              let selectedTextRange,
              selectedTextRange.isEmpty,
              let wholeDocumentRange = self.textRange(from: self.beginningOfDocument, to: self.endOfDocument),
              let text = self.text(in: wholeDocumentRange),
              !text.isEmpty else {
            return nil
        }
        
        let selectedRangeStartUTF16 = self.offset(from: self.beginningOfDocument, to: selectedTextRange.start)
        let cursorIndex = text.utf16.index(text.startIndex, offsetBy: selectedRangeStartUTF16)
        
        guard cursorIndex != text.startIndex else {
            return nil
        }
        
        var lastWordStartIndex = text.index(before: cursorIndex)
        var foundFirstAtSign = false
        while true {
            let c = text[lastWordStartIndex]
            
            if !isPermittedForAutocomplete(c) {
                if foundFirstAtSign {
                    if c != "@" {
                        // move the index forward by 1, so that the first char of the substring is the 1st @ instead of whatever comes before it
                        lastWordStartIndex = text.index(after: lastWordStartIndex)
                    }
                    break
                } else {
                    if c == "@" {
                        foundFirstAtSign = true
                    } else if c != "." {
                        // periods are allowed for domain names in mentions
                        break
                    }
                }
            }
            
            guard lastWordStartIndex > text.startIndex else {
                break
            }
            
            lastWordStartIndex = text.index(before: lastWordStartIndex)
        }
        
        return (lastWordStartIndex, foundFirstAtSign)
    }
}

enum AutocompleteState: Equatable {
    case mention(String)
    case emoji(String)
    case hashtag(String)
}

struct AutocompleteModes: OptionSet {
    static let mentions = AutocompleteModes(rawValue: 1 << 0)
    static let hashtags = AutocompleteModes(rawValue: 1 << 2)
    static let emojis   = AutocompleteModes(rawValue: 1 << 3)
    
    static let all: AutocompleteModes = [
        .mentions,
        .hashtags,
        .emojis,
    ]
    
    let rawValue: Int
    
    var triggerChars: [Character] {
        var chars: [Character] = []
        if contains(.mentions) {
            chars.append("@")
        }
        if contains(.hashtags) {
            chars.append("#")
        }
        if contains(.emojis) {
            chars.append(":")
        }
        return chars
    }
}

private func isPermittedForAutocomplete(_ c: Character) -> Bool {
    return (c >= "a" && c <= "z") || (c >= "A" && c <= "Z") || (c >= "0" && c <= "9") || c == "_"
}