import Foundation
import PlaygroundSupport

PlaygroundPage.current.needsIndefiniteExecution = true

let json = """
    [
        {
            "a": 1,
            "b": {"c": false, "d": "foo"}
        },
        {
            "a": 2,
            "b": {"c": true, "d": "[{bar}]"}
        },
    ]
    """
let jsonData = json.data(using: .utf8)!

struct Model: Codable {
    struct Inner: Codable {
        var c: Bool
        var d: String
    }
    
    var a: Int
    var b: Inner 
}

let value = try! JSONDecoder().decode([Model].self, from: jsonData)
print("Serial parsing:")
print(value)

enum JSONStructureTokenizer {
    struct Token: Equatable, CustomStringConvertible {
        enum Kind: Character, Equatable {
            case openArray = "["
            case closeArray = "]"
            case openDictionary = "{"
            case closeDictionary = "}"
            
            case stringDelimiter = "\""
        }
        
        let kind: Kind
        let index: String.UnicodeScalarView.Index
        
        var description: String {
            return "\"\(kind.rawValue)\" at \(index)"
        }
    }
    
    static func tokenize(json: String) -> [Token] {
        var tokens: [Token] = []
        
        var insideString = false
        
        for scalarIndex in json.unicodeScalars.indices {
            let scalar = json.unicodeScalars[scalarIndex]
            if let tokenKind = Token.Kind(rawValue: Character(scalar)) {    
                switch tokenKind {
                // Note: this is incredibly naive and doesn't handle escaped quotes inside strings
                case .stringDelimiter:
                    insideString.toggle()
                default:
                    if !insideString {
                        tokens.append(.init(kind: tokenKind, index: scalarIndex))
                    }
                }
            }
        }
        
        return tokens
    }
}

struct AsyncJSONDecoder<T: Decodable>: AsyncSequence {
    typealias Element = T
    
    private let json: String
    
    init(json: String) {
        self.json = json
    }
    
    func makeAsyncIterator() -> AsyncJSONIterator {
        return AsyncJSONIterator(json: json)
    }
    
    struct AsyncJSONIterator: AsyncIteratorProtocol {
        // TODO: Add better metadata to errors
        enum Error: Swift.Error {
            case invalidArray
            case invalidDictionary
        }
        
        private var lastTokenIndex: Int?
        
        private let json: String
        
        private let tokens: [JSONStructureTokenizer.Token]
        
        init(json: String) {
            self.json = json
            self.tokens = JSONStructureTokenizer.tokenize(json: json)
        }
        
        private mutating func consumeNextDictionary() throws -> (openDictionaryTokenIndex: Int, closeDictionaryTokenIndex: Int)? {
            var dictionaryNestLevel = 0
            
            var openTokenIndex: Int?
            var closeTokenIndex: Int?
            
            let remainingTokens = tokens.enumerated().dropFirst((lastTokenIndex ?? 0) + 1).dropLast()
            
            for (index, token) in remainingTokens {
                switch token.kind {
                case .openDictionary:
                    if dictionaryNestLevel == 0 && openTokenIndex == nil {    
                        openTokenIndex = index
                    } else {
                        dictionaryNestLevel += 1
                    }
                case .closeDictionary:
                    if dictionaryNestLevel == 0 {
                        closeTokenIndex = index
                        break
                    } else {
                        dictionaryNestLevel -= 1
                    }
                default: 
                    // Note: This is also very naive, and will fail in different cases of malformed JSON
                    break
                }
            }
            
            guard let openTokenIndex = openTokenIndex else {
                // No more dictionaries
                return nil
            }
            
            guard let closeTokenIndex = closeTokenIndex else {
                throw Error.invalidDictionary
            }
            
            lastTokenIndex = closeTokenIndex
            
            return (openDictionaryTokenIndex: openTokenIndex, closeDictionaryTokenIndex: closeTokenIndex)
        }
        
        private let jsonDecoder = JSONDecoder()
        
        mutating func next() async throws -> T? {
            if lastTokenIndex == nil {
                guard !tokens.isEmpty else { return nil }
                
                guard tokens.first?.kind == .openArray
                && tokens.last?.kind == .closeArray else {
                    throw Error.invalidArray
                }
            }
            
            if let nextDictionary = try consumeNextDictionary() {
                let range =  tokens[nextDictionary.openDictionaryTokenIndex].index...tokens[nextDictionary.closeDictionaryTokenIndex].index
                
                let subJSON = json[range].data(using: .utf8)!
                
                let nextValue = try jsonDecoder.decode(T.self, from: subJSON)
                
                return nextValue
            } else {
                return nil
            }
            
        }
    }
}

let asyncDecoder = AsyncJSONDecoder<Model>(json: json)

async {
    do {
        for try await value in asyncDecoder {
            print(value)
        }
        
        print("No more values")
    }
    catch {
        print("Error: \(error)")
    }
}