//
//  JSONPrettyPrinter.swift
//  MongoView
//
//  Created by Shadowfacts on 8/12/20.
//  Copyright © 2020 Shadowfacts. All rights reserved.
//

import Foundation

class JSONPrettyPrinter {
    let options: Options
    private var strings = [String]()
    private var currentIndent = ""
    
    init(options: Options) {
        self.options = options
    }
    
    func prettify(_ string: String) throws -> String {
        let object = try JSONSerialization.jsonObject(with: string.data(using: .utf8)!, options: [])
        
        doPrettify(object)

        return strings.joined(separator: "")
    }
    
    private func indent() {
        currentIndent += "\t"
    }
    
    private func outdent() {
        currentIndent = String(currentIndent.dropLast())
    }
    
    private func doPrettify(_ object: Any) {
        if let dict = object as? [String: Any] {
            if options.contains(.convertMongoObjects) && tryPrettifyMongoObject(dict) {
                return
            }
            indent()
            strings.append("{\n")
            for (index, k) in dict.keys.sorted().enumerated() {
                strings.append("\(currentIndent)\"\(escape(k))\"")
                strings.append(": ")
                doPrettify(dict[k]!)
                if index != dict.count - 1 {
                    strings.append(",\n")
                }
            }
            outdent()
            strings.append("\n\(currentIndent)}")
        } else if let arr = object as? [Any] {
            indent()
            strings.append("[\n")
            for (index, v) in arr.enumerated() {
                strings.append(currentIndent)
                doPrettify(v)
                if index != arr.count - 1 {
                    strings.append(",\n")
                }
            }
            outdent()
            strings.append("\n\(currentIndent)]")
        } else if let str = object as? String {
            strings.append("\"\(escape(str))\"")
        } else if let num = object as? NSNumber {
            strings.append(num.stringValue)
        } else if object is NSNull {
            strings.append("null")
        } else {
            fatalError("unhandled object type: \(String(describing: object))")
        }
    }
    
    private func tryPrettifyMongoObject(_ dict: [String: Any]) -> Bool {
        guard dict.count == 1 else { return false }
        if let value = dict["$oid"] as? String {
            strings.append("ObjectId(\"\(value)\")")
            return true
        }
        if let value = dict["$date"] as? String {
            strings.append("Date(\"\(value)\")")
            return true
        }
        return false
    }
    
    private func escape(_ str: String) -> String {
        var str = str
        var index = str.startIndex
        while index < str.endIndex {
            let c = str[index]
            if c == "\\" || c == "\"" {
                str.replaceSubrange(index..<str.index(after: index), with: "\\\(c)")
                index = str.index(after: index)
            } else if c == "\n" {
                str.replaceSubrange(index..<str.index(after: index), with: "\\n")
                index = str.index(after: index)
            } else if c == "\r" {
                str.replaceSubrange(index..<str.index(after: index), with: "\\r")
                index = str.index(after: index)
            }
            index = str.index(after: index)
        }
        return str
    }
}

extension JSONPrettyPrinter {
    struct Options: OptionSet {
        let rawValue: Int
        
        static let convertMongoObjects = Options(rawValue: 1 << 0)
    }
}