import Foundation
import Splash

@_cdecl("highlight_swift")
public func highlight(codePtr: UnsafePointer<UInt8>, codeLen: UInt64, htmlLenPtr: UnsafeMutablePointer<UInt64>) -> UnsafeMutablePointer<UInt8> {
    let buf = UnsafeBufferPointer(start: codePtr, count: Int(codeLen))
    let data = Data(buffer: buf)
    let code = String(data: data, encoding: .utf8)!
    
    let highligher = SyntaxHighlighter(format: MyOutputFormat())
    var html = highligher.highlight(code)
    let outPtr = UnsafeMutableBufferPointer<UInt8>.allocate(capacity: html.utf8.count)
    _ = html.withUTF8 { buf in
        buf.copyBytes(to: outPtr, count: buf.count)
    }
    htmlLenPtr.pointee = UInt64(outPtr.count)
    return outPtr.baseAddress!
}

struct MyOutputFormat: OutputFormat {
    func makeBuilder() -> Builder {
        return Builder()
    }
    
    struct Builder: OutputBuilder {
        typealias Output = String
        
        private var html = ""
        private var pendingToken: (string: String, type: TokenType)?
        private var pendingWhitespace: String?
        
        mutating func addToken(_ token: String, ofType type: TokenType) {
            if var pendingToken = pendingToken {
                guard pendingToken.type != type else {
                    pendingWhitespace.map { pendingToken.string += $0 }
                    pendingWhitespace = nil
                    pendingToken.string += token
                    self.pendingToken = pendingToken
                    return
                }
            }
            
            appendPending()
            pendingToken = (token, type)
        }
        
        mutating func addPlainText(_ text: String) {
            appendPending()
            html.append(contentsOf: text.escapingHTMLEntities())
        }
        
        mutating func addWhitespace(_ whitespace: String) {
            if pendingToken != nil {
                pendingWhitespace = (pendingWhitespace ?? "") + whitespace
            } else {
                html.append(whitespace)
            }
        }
        
        mutating func build() -> String {
            appendPending()
            return html
        }
        
        private mutating func appendPending() {
            if let pendingToken = pendingToken {
                let cls: String
                switch pendingToken.type {
                case .keyword:
                    cls = "hl-kw"
                case .string:
                    cls = "hl-str"
                case .type:
                    cls = "hl-type"
                case .call:
                    cls = "hl-fn"
                case .number:
                    cls = "hl-num"
                case .comment:
                    cls = "hl-cmt"
                case .property:
                    cls = "hl-var"
                case .dotAccess:
                    cls = "hl-prop"
                case .preprocessing:
                    cls = ""
                case .custom(_):
                    cls = ""
                }
                html.append("""
                <span class="\(cls)">\(pendingToken.string.escapingHTMLEntities())</span>
                """)
                
                self.pendingToken = nil
            }
            
            if let pendingWhitespace = pendingWhitespace {
                html.append(pendingWhitespace)
                self.pendingWhitespace = nil
            }
        }
    }
}

// copied from Splash, where it's internal
extension StringProtocol {
    func escapingHTMLEntities() -> String {
        return String(flatMap { character -> String in
            switch character {
            case "&":
                return "&amp;"
            case "<":
                return "&lt;"
            case ">":
                return "&gt;"
            default:
                return String(character)
            }
        })
    }
}