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

import AppKit

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 let 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
    
    init(text: String) {
        self.text = text
        self.attributed = NSMutableAttributedString(attributedString: NSAttributedString(string: text))
    }
    
    init(mutableAttributed: NSMutableAttributedString) {
        self.text = mutableAttributed.string
        self.attributed = mutableAttributed
    }
    
    private func print(_ str: @autoclosure () -> String) {
        if debug {
            Swift.print("\(indent)\(str())")
        }
    }
    
    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)]
    }
    
    @discardableResult
    private func consume() -> Unicode.Scalar? {
        let c = peek()
        currentIndex = text.index(after: currentIndex)
        return c
    }
    
    func 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()
    }
    
    private func emit(token: TokenType, range: NSRange) {
        let color: NSColor
        switch token {
        case .string:
            color = .systemRed
        case .number:
            color = .systemBlue
        case .punctuation:
            color = .systemTeal
        case .identifier:
            return
        }
        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 startChar = consume()
        while currentIndex < text.endIndex && peek() != startChar && (currentIndex == text.startIndex || text[text.index(before: currentIndex)] != "\\") {
            consume()
        }
        if currentIndex < text.endIndex {
            consume() // string closing quote
        }
        print("String: \(text[stringStart..<currentIndex])")
        emit(token: .string, 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 currentIndex < text.endIndex {
            if peek() == "$" {
                emitTemplateStringFragment()
                consume() // $
                if peek() == "{" {
                    consume() // {
                    emitTemplateStringFragment()
                    consumeTemplateStringExpression()
                    stringFragmentStart = currentIndex
                    if currentIndex < text.endIndex && peek() == "}" {
                        consume() // }
                    }

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

    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()
        }
        indent = String(indent.dropLast(2))
        if currentIndex < text.endIndex {
            consume() // )
            print("Closing )")
            emit(token: .punctuation, range: prevCharRange())
        }
    }
    
    private func consumeObject() {
        consume() // {
        print("Opening {")
        emit(token: .punctuation, range: prevCharRange())
        indent += "  "
        object:
        while currentIndex < text.endIndex && peek() != "}" {
            consumeObjectKey()
            if peek() == ":" {
                consume() // :
                emit(token: .punctuation, range: prevCharRange())
                consumeWhitespace()
                while currentIndex < text.endIndex && peek() != "," && peek() != "}" {
                    consumeExpression()
                }
            }
            consumeWhitespace()

            if peek() == "," {
                consume() // ,
                emit(token: .punctuation, range: prevCharRange())
                continue
            } else {
                break
            }
        }
        indent = String(indent.dropLast(2))
        if currentIndex < text.endIndex {
            consume() // }
            print("Closing }")
            emit(token: .punctuation, range: prevCharRange())
        }
    }
    
    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 += "  "
        array:
        while currentIndex < text.endIndex {
            consumeWhitespace()
            if peek() == "," {
                consume() // ,
                print("Array separator")
                emit(token: .punctuation, range: prevCharRange())
            } else if peek() == "]" {
                break array
            } else {
                print("Array element")
                while let char = peek(), char != ",", char != "]" {
                    consumeExpression()
                }
            }
        }
        indent = String(indent.dropLast(2))
        if currentIndex < text.endIndex {
            consume() // ]
            print("Closing ]")
            emit(token: .punctuation, range: prevCharRange())
        }
    }
    
    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()
        }
        indent = String(indent.dropLast(2))
    }
    
}

extension JavaScriptHighlighter {
    enum TokenType {
        case identifier
        case punctuation
        case number
        case string
    }
}