//
//  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: "+-*/<>=")

class JavaScriptHighlighter {
    private let text: String
    private var attributed: NSMutableAttributedString!
    private var currentIndex: String.Index!
    private var indent = ""
    
    init(text: String) {
        self.text = text
    }
    
    private func print(_ str: String) {
        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(mutableAttributed: NSMutableAttributedString? = nil) -> NSAttributedString {
        attributed = mutableAttributed ?? NSMutableAttributedString(attributedString: NSAttributedString(string: text))
        
        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 currentIndex < text.endIndex {
            consumeExpression()
        }

        return attributed
    }
    
    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)
    }

    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 {
            consume()
        }
    }
    
    private func consumeWhitespace(newlines: Bool = false) {
        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
        consume() // opening `
        while currentIndex < text.endIndex {
            if peek(length: 2) == "${" {
                consume() // $
                consume() // {
                print("Template string fragment: '\(text[stringFragmentStart!..<currentIndex])'")
                emit(token: .string, range: range(from: stringFragmentStart!, to: currentIndex))
                consumeTemplateStringExpression()
                stringFragmentStart = currentIndex
                if currentIndex < text.endIndex && peek() == "}" {
                    consume()
                }
            } else if peek() == "`" {
                stringFragmentStart = stringFragmentStart ?? currentIndex
                consume() // `
                print("Template string fragment: '\(text[stringFragmentStart!..<currentIndex])'")
                emit(token: .string, range: range(from: stringFragmentStart!, to: currentIndex))
                stringFragmentStart = nil
                break
            } else {
                consume()
            }
        }
        if let start = stringFragmentStart {
            print("Template string fragment: '\(text[start..<currentIndex])'")
            emit(token: .string, range: range(from: start, to: currentIndex))
        }
    }
    
    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() // :
                while currentIndex < text.endIndex && peek() != "," && peek() != "}" {
                    consumeExpression()
                }
                if currentIndex < text.endIndex {
                    break object
                }
            } else if peek() == "," {
                continue
            } else {
                break
            }
        }
        indent = String(indent.dropLast(2))
        if currentIndex < text.endIndex {
            consume() // }
            print("Closing }")
            emit(token: .punctuation, range: prevCharRange())
        }
    }
    
    private func consumeObjectKey() {
        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])'")
    }
    
    private func consumeDotLookup() {
        consume() // .
        print("Dot lookup")
        emit(token: .punctuation, range: prevCharRange())
    }
    
    private func consumeArray() {
        consume() // [
        print("Opening [")
        emit(token: .punctuation, range: prevCharRange())
        indent += "  "
        array:
        while currentIndex < text.endIndex && peek() != "]" {
            print("Array element")
            while currentIndex < text.endIndex {
                if peek() == "," {
                    consume() // ,
                    print("Array separator")
                    emit(token: .punctuation, range: prevCharRange())
                    break
                } else if peek() == "]" {
                    break array
                } else {
                    indent += "  "
                    consumeExpression()
                    indent = String(indent.dropLast(2))
                }
            }
        }
        indent = String(indent.dropLast(2))
        if currentIndex < text.endIndex {
            consume() // ]
            print("Closing ]")
            emit(token: .punctuation, range: prevCharRange())
        }
    }
    
}

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