enum JSON {
  	case array([JSON])
  	case object([String: JSON])
  	case number(Int)
  	case string(String)
  	case bool(Bool)
}

extension JSON: ExpressibleByIntegerLiteral {
	init(integerLiteral: Int) {
		self = .number(integerLiteral)
	}
}

extension JSON: ExpressibleByBooleanLiteral {
	init(booleanLiteral: Bool) {
		self = .bool(booleanLiteral)
	}
}

extension JSON: ExpressibleByStringLiteral {
	typealias StringLiteralType = String
	typealias ExtendedGraphemeClusterLiteralType = String
	typealias UnicodeScalarLiteralType = String
	
	init(stringLiteral: String) {
		self = .string(stringLiteral)
	}
	
	init(extendedGraphemeClusterLiteral: String) {
		self.init(stringLiteral: extendedGraphemeClusterLiteral)
	}
	
	init(unicodeScalarLiteral: String) {
		self.init(stringLiteral: unicodeScalarLiteral)
	}
}

extension JSON: ExpressibleByDictionaryLiteral {
	typealias Key = String
	typealias Value = JSON
	
	init(dictionaryLiteral: (Key, Value)...) {
		var dict: [Key: Value] = [:]
		for (key, value) in dictionaryLiteral {
			dict[key] = value
		}
		self = .object(dict)
	}

}

extension JSON: ExpressibleByArrayLiteral {
	init(arrayLiteral: Value...) {
		self = .array(arrayLiteral)
	}
}

func array(_ input: JSON) -> JSON? {
  switch input {
  case .array: return input
  default: return nil
  }
}

func object(_ input: JSON) -> JSON? {
  switch input {
  case .object: return input
  default: return nil
  }
}

func key(_ key: String) -> (JSON) -> JSON? {
  return { input in
	switch object(input) {
		case let .some(.object(dict)):
			return dict[key]
		default:
			return nil
	}
  }
}

func index(_ index: Int) -> (JSON) -> JSON? {
  return { input in
	switch array(input) {
		case let .some(.array(array)):
			if index > array.count - 1 {
				return nil
			}
			else {
				return array[index]
			}			
		default:
			return nil
	}
  }
}

func number(_ input: JSON) -> JSON? {
	switch input {
		case .number: return input
		default: return nil
	}
}

func bool(_ input: JSON) -> JSON? {
	switch input {
		case .bool: return input
		default: return nil
	}
}

func string(_ input: JSON) -> JSON? {
	switch input {
		case .string: return input
		default: return nil
	}
}

typealias JSONQuery = (JSON) -> JSON?

infix operator >>: AdditionPrecedence
func >> (lhs: @escaping JSONQuery, rhs: @escaping JSONQuery) -> JSONQuery {
  return { input in 
    guard let result = lhs(input) else { return nil }
    return rhs(result)
  }
}

let json: JSON = ["herp": [1, 2, 4]]
let query = object >> key("herp") >> array >> index(2) >> number
print(query(json) ?? "booo")