//
//  JavaScriptHighlighter.swift
//  MongoView
//
//  Created by Shadowfacts on 4/1/20.
//  Copyright © 2020 Shadowfacts. All rights reserved.
//

import AppKit
import OSLog

fileprivate let log = OSLog(subsystem: "space.vaccor.MongoView.JavaScriptHighlighter", category: .pointsOfInterest)

fileprivate let identifiers: CharacterSet = {
    var set = CharacterSet.alphanumerics
    set.insert(charactersIn: "$_")
    return set
}()
fileprivate let identifierStarts: CharacterSet = {
    var set = identifiers
    set.subtract(.decimalDigits)
    return set
}()
fileprivate let operators = CharacterSet(charactersIn: "+-*/<>=")
fileprivate let expressionEnds = CharacterSet(charactersIn: ",]});")

class JavaScriptHighlighter {
    private var text: String!
    private var attributed: NSMutableAttributedString!
    private var currentIndex: String.Index!
    private var _indent = ""
    private(set) var tokens = [(token: TokenType, range: NSRange)]()
    var debug = false

    private func print(_ str: @autoclosure () -> String) {
        #if DEBUG
        if debug {
            Swift.print("\(_indent)\(str())")
        }
        #endif
    }
    
    private func indent() {
        #if DEBUG
        _indent += "  "
        #endif
    }
    private func outdent() {
        #if DEBUG
        _indent = String(_indent.dropLast(2))
        #endif
    }

    private func range(from: String.Index, to: String.Index) -> NSRange {
        return NSRange(from..<to, in: text)
    }
    
    private func prevCharRange() -> NSRange {
        return range(from: text.index(before: currentIndex), to: currentIndex)
    }
    
    private func peek() -> Unicode.Scalar? {
        guard currentIndex < text.endIndex else {
            return nil
        }
        return text.unicodeScalars[currentIndex]
    }
    
    private func peek(length: Int) -> Substring {
        let realLength = min(length, text.distance(from: currentIndex, to: text.endIndex))
        return text[currentIndex..<text.index(currentIndex, offsetBy: realLength)]
    }
    
    private func consume() {
        currentIndex = text.index(after: currentIndex)
    }
    
    private func getAndConsume() -> Unicode.Scalar? {
        let c = peek()
        consume()
        return c
    }
    
    func highlight(text: String) {
        self.highlight(attributed: NSMutableAttributedString(attributedString: NSAttributedString(string: text)))
    }
    
    func highlight(attributed: NSMutableAttributedString) {
        self.attributed = attributed
        self.text = attributed.string
        self.tokens = []
        
        os_signpost(.begin, log: log, name: "highlight")
        defer { os_signpost(.end, log: log, name: "highlight") }
        
        attributed.beginEditing()
        let fullRange = NSRange(location: 0, length: attributed.length)
        attributed.setAttributes([
            .foregroundColor: NSColor.textColor,
            .font: NSFont.monospacedSystemFont(ofSize: 13, weight: .regular)
        ], range: fullRange)

        currentIndex = text.startIndex
        while let char = peek(), !expressionEnds.contains(char) {
            consumeExpression()
        }

        attributed.endEditing()
    }
    
    func inserted(character: Unicode.Scalar, index: Int) -> Bool {
        os_signpost(.begin, log: log, name: "inserted(character:index:)")
        defer { os_signpost(.end, log: log, name: "inserted(character:index:)") }

        if let (token, range) = tokenAndRange(fullyContaining: index) {
            switch token {
            case .identifier:
                let set = range.location == index ? identifierStarts : identifiers
                guard set.contains(character) else { return false }
            case let .string(stringChar):
                guard character != stringChar || (index > 0 && Unicode.Scalar((text as NSString).character(at: index - 1)) == "\\"),
                    character != "\\" else { return false }
            case .number:
                if character == "." && (attributed.string as NSString).substring(with: range).contains(Character(".")) {
                    return false
                }
            default:
                return false
            }
            var attributes: [NSAttributedString.Key: Any] = [
                .font: NSFont.monospacedSystemFont(ofSize: 13, weight: .regular)
            ]
            if let color = token.color {
                attributes[.foregroundColor] = color
            }
            attributed.addAttributes(attributes, range: NSRange(location: index, length: 1))
            tokens = tokens.map { (token, range) in
                if range.contains(index) {
                    return (token, NSRange(location: range.location, length: range.length + 1))
                } else if range.location >= index {
                    return (token, NSRange(location: range.location + 1, length: range.length))
                } else {
                    return (token, range)
                }
            }
            return true
        }
        return false
    }
    
    func removed(index: Int) -> Bool {
        os_signpost(.begin, log: log, name: "removed(at:)")
        defer { os_signpost(.end, log: log, name: "removed(at:)") }

        if let (token, range) = tokenAndRange(fullyContaining: index) {
            switch token {
            case .string(_):
                guard index > range.location && index < range.location + range.length - 1 else { return false }
                if index > 0 && Unicode.Scalar((text as NSString).character(at: index - 1)) == "\\" {
                    // removed character
                    return false
                }
            case .number:
                break
            default:
                return false
            }
            tokens = tokens.map { (token, range) in
                if range.contains(index) {
                    return (token, NSRange(location: range.location, length: range.length - 1))
                } else if range.location >= index {
                    return (token, NSRange(location: range.location - 1, length: range.length))
                } else {
                    return (token, range)
                }
            }
            return true
        }
        return false
    }
    
    func token(at index: Int) -> TokenType? {
        for (token, range) in tokens where range.contains(index) {
            return token
        }
        return nil
    }
    
    private func tokenAndRange(fullyContaining index: Int) -> (TokenType, NSRange)? {
        for (token, range) in tokens where index > range.location && index < range.location + range.length {
            return (token, range)
        }
        return nil
    }
    
    private func emit(token: TokenType, range: NSRange) {
        if let color = token.color {
            attributed.addAttribute(.foregroundColor, value: color, range: range)
        }
        tokens.append((token, range))
    }

    private func consumeExpression() {
        consumeWhitespace()
        
        guard let char = peek() else { return }
        
        if identifierStarts.contains(char) {
            consumeIdentifier()
        } else if char == "'" || char == "\"" {
            consumeString()
        } else if char == "`" {
            consumeTemplateString()
        } else if CharacterSet.decimalDigits.contains(char) {
            consumeNumber()
        } else if operators.contains(char) {
            consumeOperator()
        } else if char == "(" {
            consumeFunctionCallOrGrouping()
        } else if char == "{" {
            consumeObject()
        } else if char == "[" {
            consumeArray()
        } else if char == "." {
            consumeDotLookup()
        } else if char == "?" {
            consumeTernaryExpression()
        } else {
            consume()
        }
        
        consumeWhitespace()
    }
    
    private func consumeWhitespace(newlines: Bool = true) {
        let charSet = newlines ? CharacterSet.whitespacesAndNewlines : .whitespaces
        while let char = peek(), charSet.contains(char) {
            consume()
        }
    }

    private func consumeIdentifier() {
        let identifierStart = currentIndex!
        while let char = peek(), identifiers.contains(char) {
            consume()
        }
        print("Identifier: '\(text[identifierStart..<currentIndex])'")
    }

    private func consumeNumber() {
        let numberStart = currentIndex!
        while let char = peek(), CharacterSet.decimalDigits.contains(char) {
            consume()
        }
        if currentIndex < text.endIndex && peek() == "." {
            consume()
            while let char = peek(), CharacterSet.decimalDigits.contains(char) {
                consume()
            }
        }
        print("Number: \(text[numberStart..<currentIndex])")
        emit(token: .number, range: range(from: numberStart, to: currentIndex))
    }
    
    private func consumeString() {
        let stringStart = currentIndex!
        let quote = peek()!
        consume() // opening quote
        var prevChar: Unicode.Scalar?
        var char = peek()
        while currentIndex < text.endIndex {
            consume()
            if char == quote && prevChar != "\\" {
                break
            }
            prevChar = char
            char = peek()
        }
        print("String: \(text[stringStart..<currentIndex])")
        emit(token: .string(quote), range: range(from: stringStart, to: currentIndex))
    }
    
    private func consumeTemplateString() {
        var stringFragmentStart: String.Index? = currentIndex
        
        func emitTemplateStringFragment() {
            guard stringFragmentStart != currentIndex else { return }
            print("Template string fragment: '\(text[stringFragmentStart!..<currentIndex])'")
            emit(token: .string("`"), range: range(from: stringFragmentStart!, to: currentIndex))
            stringFragmentStart = currentIndex
        }
        
        consume() // opening `
        
        while let char = peek() {
            if char == "$" {
                emitTemplateStringFragment()
                consume() // $
                if peek() == "{" {
                    consume() // {
                    emitTemplateStringFragment()
                    consumeTemplateStringExpression()
                    stringFragmentStart = currentIndex
                    if currentIndex < text.endIndex && peek() == "}" {
                        consume() // }
                    }

                }
            } else if char == "`" {
                stringFragmentStart = stringFragmentStart ?? currentIndex
                consume() // `
                emitTemplateStringFragment()
                stringFragmentStart = nil
                break
            } else {
                consume()
            }
        }
        if stringFragmentStart != nil {
            emitTemplateStringFragment()
        }
    }
    
    private func consumeTemplateStringExpression() {
        indent()
        while currentIndex < text.endIndex && peek() != "}" {
            consumeExpression()
        }
        outdent()
    }

    private func consumeOperator() {
        print("Operator: \(peek()!)")
        consume()
    }
    
    private func consumeFunctionCallOrGrouping() {
        consume() // (
        print("Opening (")
        emit(token: .punctuation, range: prevCharRange())
        indent()
        while currentIndex < text.endIndex && peek() != ")" {
            consumeExpression()
        }
        outdent()
        if currentIndex < text.endIndex {
            consume() // )
            print("Closing )")
            emit(token: .punctuation, range: prevCharRange())
        }
    }
    
    private func consumeObject() {
        consume() // {
        print("Opening {")
        emit(token: .punctuation, range: prevCharRange())
        indent()
        while let char = peek() {
            if char == "}" {
                consume() // }
                print("Closing }")
                emit(token: .punctuation, range: prevCharRange())
                break
            } else if char == "'" || char == "\"" {
                let keyStart = currentIndex!
                consumeString()
                print("Object key: '\(text[keyStart..<currentIndex])'")
            } else if identifierStarts.contains(char) {
                let keyStart = currentIndex!
                consumeIdentifier()
                print("Object key: '\(text[keyStart..<currentIndex])'")
            } else if char == ":" || char == "," {
                consume() // :
                emit(token: .punctuation, range: prevCharRange())
            } else {
                consumeExpression()
            }
        }
        outdent()
    }
    
    private func consumeObjectKey() {
        consumeWhitespace()
        guard let char = peek() else { return }
        let keyStart = currentIndex!
        if identifierStarts.contains(char) {
            consumeIdentifier()
        } else if char == "'" || char == "\"" {
            consumeString()
        }
        print("Object key: '\(text[keyStart..<currentIndex])'")
        consumeWhitespace()
    }
    
    private func consumeDotLookup() {
        consume() // .
        guard let char = peek() else { return }
        if identifierStarts.contains(char) {
            print("Dot lookup")
            emit(token: .punctuation, range: prevCharRange())
            consumeIdentifier()
        } else if CharacterSet.decimalDigits.contains(char) {
            let numberStart = text.index(before: currentIndex)
            while let char = peek(), CharacterSet.decimalDigits.contains(char) {
                consume()
            }
            print("Number: \(text[numberStart..<currentIndex])")
            emit(token: .number, range: range(from: numberStart, to: currentIndex))
        }
    }
    
    private func consumeArray() {
        consume() // [
        print("Opening [")
        emit(token: .punctuation, range: prevCharRange())
        indent()
        while let char = peek() {
            if char == "]" {
                consume() // ]
                print("Closing ]")
                emit(token: .punctuation, range: prevCharRange())
                break
            } else if char == "," {
                consume() // ,
                print("Array separator")
                emit(token: .punctuation, range: prevCharRange())
            } else {
                consumeExpression()
            }
        }
        outdent()
    }
    
    func consumeTernaryExpression() {
        consume() // ?
        print("Ternary expression")
        indent()
        print("Ternary true result")
        while let char = peek(), char != ":" {
            consumeExpression() // true result
        }
        consume() // :
        print("Ternary false result")
        while let char = peek(), !expressionEnds.contains(char) {
            consumeExpression()
        }
        outdent()
    }
    
}

extension JavaScriptHighlighter {
    enum TokenType {
        case identifier
        case punctuation
        case number
        case string(Unicode.Scalar)
        
        var color: NSColor? {
            switch self {
            case .string(_):
                return .systemRed
            case .number:
                return .systemBlue
            case .punctuation:
                return .systemTeal
            case .identifier:
                return nil
            }
        }
    }
}