// Sources/SwiftProtobuf/AnyMessageStorage.swift - Custom storage for Any WKT

//

// Copyright (c) 2014 - 2017 Apple Inc. and the project authors

// Licensed under Apache License v2.0 with Runtime Library Exception

//

// See LICENSE.txt for license information:

// https://github.com/apple/swift-protobuf/blob/main/LICENSE.txt

//

// -----------------------------------------------------------------------------

///

/// Hand written storage class for Google_Protobuf_Any to support on demand

/// transforms between the formats.

///

// -----------------------------------------------------------------------------

#if canImport(FoundationEssentials)

import FoundationEssentials

#else

import Foundation

#endif

private func serializeAnyJSON(

    for message: any Message,

    typeURL: String,

    options: JSONEncodingOptions

) throws -> String {

    var visitor = try JSONEncodingVisitor(type: type(of: message), options: options)

    visitor.startObject(message: message)

    visitor.encodeField(name: "@type", stringValue: typeURL)

    if let m = message as? (any _CustomJSONCodable) {

        let value = try m.encodedJSONString(options: options)

        visitor.encodeField(name: "value", jsonText: value)

    } else {

        try message.traverse(visitor: &visitor)

    }

    visitor.endObject()

    return visitor.stringResult

}

private func emitVerboseTextForm(visitor: inout TextFormatEncodingVisitor, message: any Message, typeURL: String) {

    let url: String

    if typeURL.isEmpty {

        url = buildTypeURL(forMessage: message, typePrefix: defaultAnyTypeURLPrefix)

    } else {

        url = typeURL

    }

    visitor.visitAnyVerbose(value: message, typeURL: url)

}

private func asJSONObject(body: [UInt8]) -> Data {

    let asciiOpenCurlyBracket = UInt8(ascii: "{")

    let asciiCloseCurlyBracket = UInt8(ascii: "}")

    var result = Data()

    result.reserveCapacity(body.count + 2)

    result.append(asciiOpenCurlyBracket)

    result.append(contentsOf: body)

    result.append(asciiCloseCurlyBracket)

    return result

}

private func unpack(

    contentJSON: [UInt8],

    extensions: any ExtensionMap,

    options: JSONDecodingOptions,

    as messageType: any Message.Type

) throws -> any Message {

    guard messageType is any _CustomJSONCodable.Type else {

        let contentJSONAsObject = asJSONObject(body: contentJSON)

        return try messageType.init(jsonUTF8Bytes: contentJSONAsObject, extensions: extensions, options: options)

    }

    var value = String()

    try contentJSON.withUnsafeBytes { (body: UnsafeRawBufferPointer) in

        if body.count > 0 {

            // contentJSON will be the valid JSON for inside an object (everything but

            // the '{' and '}', so minimal validation is needed.

            var scanner = JSONScanner(source: body, options: options, extensions: extensions)

            while !scanner.complete {

                let key = try scanner.nextQuotedString()

                try scanner.skipRequiredColon()

                if key == "value" {

                    value = try scanner.skip()

                    break

                }

                if !options.ignoreUnknownFields {

                    // The only thing within a WKT should be "value".

                    throw AnyUnpackError.malformedWellKnownTypeJSON

                }

                let _ = try scanner.skip()

                try scanner.skipRequiredComma()

            }

            if !options.ignoreUnknownFields && !scanner.complete {

                // If that wasn't the end, then there was another key, and WKTs should

                // only have the one when not skipping unknowns.

                throw AnyUnpackError.malformedWellKnownTypeJSON

            }

        }

    }

    return try messageType.init(jsonString: value, extensions: extensions, options: options)

}

internal class AnyMessageStorage {

    // The two properties generated Google_Protobuf_Any will reference.

    var _typeURL = String()

    var _value: Data {

        // Remapped to the internal `state`.

        get {

            switch state {

            case .binary(let value):

                return value

            case .message(let message):

                do {

                    return try message.serializedBytes(partial: true)

                } catch {

                    return Data()

                }

            case .contentJSON(let contentJSON, let options):

                guard let messageType = Google_Protobuf_Any.messageType(forTypeURL: _typeURL) else {

                    return Data()

                }

                do {

                    let m = try unpack(

                        contentJSON: contentJSON,

                        extensions: SimpleExtensionMap(),

                        options: options,

                        as: messageType

                    )

                    return try m.serializedBytes(partial: true)

                } catch {

                    return Data()

                }

            }

        }

        set {

            state = .binary(newValue)

        }

    }

    enum InternalState {

        // a serialized binary

        // Note: Unlike contentJSON below, binary does not bother to capture the

        // decoding options. This is because the actual binary format is the binary

        // blob, i.e. - when decoding from binary, the spec doesn't include decoding

        // the binary blob, it is pass through. Instead there is a public api for

        // unpacking that takes new options when a developer decides to decode it.

        case binary(Data)

        // a message

        case message(any Message)

        // parsed JSON with the @type removed and the decoding options.

        case contentJSON([UInt8], JSONDecodingOptions)

    }

    var state: InternalState = .binary(Data())

    // This property is used as the initial default value for new instances of the type.

    // The type itself is protecting the reference to its storage via CoW semantics.

    // This will force a copy to be made of this reference when the first mutation occurs;

    // hence, it is safe to mark this as `nonisolated(unsafe)`.

    static nonisolated(unsafe) let defaultInstance = AnyMessageStorage()

    private init() {}

    init(copying source: AnyMessageStorage) {

        _typeURL = source._typeURL

        state = source.state

    }

    func isA<M: Message>(_ type: M.Type) -> Bool {

        if _typeURL.isEmpty {

            return false

        }

        let encodedType = typeName(fromURL: _typeURL)

        return encodedType == M.protoMessageName

    }

    // This is only ever called with the expectation that target will be fully

    // replaced during the unpacking and never as a merge.

    func unpackTo<M: Message>(

        target: inout M,

        extensions: (any ExtensionMap)?,

        options: BinaryDecodingOptions

    ) throws {

        guard isA(M.self) else {

            throw AnyUnpackError.typeMismatch

        }

        switch state {

        case .binary(let data):

            target = try M(serializedBytes: data, extensions: extensions, partial: true, options: options)

        case .message(let msg):

            if let message = msg as? M {

                // Already right type, copy it over.

                target = message

            } else {

                // Different type, serialize and parse.

                let bytes: [UInt8] = try msg.serializedBytes(partial: true)

                target = try M(serializedBytes: bytes, extensions: extensions, partial: true)

            }

        case .contentJSON(let contentJSON, let options):

            target =

                try unpack(

                    contentJSON: contentJSON,

                    extensions: extensions ?? SimpleExtensionMap(),

                    options: options,

                    as: M.self

                ) as! M

        }

    }

    // Called before the message is traversed to do any error preflights.

    // Since traverse() will use _value, this is our chance to throw

    // when _value can't.

    func preTraverse() throws {

        switch state {

        case .binary:

            // Nothing to be checked.

            break

        case .message:

            // When set from a developer provided message, partial support

            // is done. Any message that comes in from another format isn't

            // checked, and transcoding the isInitialized requirement is

            // never inserted.

            break

        case .contentJSON(let contentJSON, let options):

            // contentJSON requires we have the type available for decoding.

            guard let messageType = Google_Protobuf_Any.messageType(forTypeURL: _typeURL) else {

                throw BinaryEncodingError.anyTranscodeFailure

            }

            do {

                // Decodes the full JSON and then discard the result.

                // The regular traversal will decode this again by querying the

                // `value` field, but that has no way to fail.  As a result,

                // we need this to accurately handle decode errors.

                _ = try unpack(

                    contentJSON: contentJSON,

                    extensions: SimpleExtensionMap(),

                    options: options,

                    as: messageType

                )

            } catch {

                throw BinaryEncodingError.anyTranscodeFailure

            }

        }

    }

}

/// Custom handling for Text format.

extension AnyMessageStorage {

    func decodeTextFormat(typeURL url: String, decoder: inout TextFormatDecoder) throws {

        // Decoding the verbose form requires knowing the type.

        _typeURL = url

        guard let messageType = Google_Protobuf_Any.messageType(forTypeURL: url) else {

            // The type wasn't registered, can't parse it.

            throw TextFormatDecodingError.malformedText

        }

        let terminator = try decoder.scanner.skipObjectStart()

        var subDecoder = try TextFormatDecoder(

            messageType: messageType,

            scanner: decoder.scanner,

            terminator: terminator

        )

        if messageType == Google_Protobuf_Any.self {

            var any = Google_Protobuf_Any()

            try any.decodeTextFormat(decoder: &subDecoder)

            state = .message(any)

        } else {

            var m = messageType.init()

            try m.decodeMessage(decoder: &subDecoder)

            state = .message(m)

        }

        decoder.scanner = subDecoder.scanner

        if try decoder.nextFieldNumber() != nil {

            // Verbose any can never have additional keys.

            throw TextFormatDecodingError.malformedText

        }

    }

    // Specialized traverse for writing out a Text form of the Any.

    // This prefers the more-legible "verbose" format if it can

    // use it, otherwise will fall back to simpler forms.

    internal func textTraverse(visitor: inout TextFormatEncodingVisitor) {

        switch state {

        case .binary(let valueData):

            if let messageType = Google_Protobuf_Any.messageType(forTypeURL: _typeURL) {

                // If we can decode it, we can write the readable verbose form:

                do {

                    let m = try messageType.init(serializedBytes: valueData, partial: true)

                    emitVerboseTextForm(visitor: &visitor, message: m, typeURL: _typeURL)

                    return

                } catch {

                    // Fall through to just print the type and raw binary data.

                }

            }

            if !_typeURL.isEmpty {

                try! visitor.visitSingularStringField(value: _typeURL, fieldNumber: 1)

            }

            if !valueData.isEmpty {

                try! visitor.visitSingularBytesField(value: valueData, fieldNumber: 2)

            }

        case .message(let msg):

            emitVerboseTextForm(visitor: &visitor, message: msg, typeURL: _typeURL)

        case .contentJSON(let contentJSON, let options):

            // If we can decode it, we can write the readable verbose form:

            if let messageType = Google_Protobuf_Any.messageType(forTypeURL: _typeURL) {

                do {

                    let m = try unpack(

                        contentJSON: contentJSON,

                        extensions: SimpleExtensionMap(),

                        options: options,

                        as: messageType

                    )

                    emitVerboseTextForm(visitor: &visitor, message: m, typeURL: _typeURL)

                    return

                } catch {

                    // Fall through to just print the raw JSON data

                }

            }

            if !_typeURL.isEmpty {

                try! visitor.visitSingularStringField(value: _typeURL, fieldNumber: 1)

            }

            // Build a readable form of the JSON:

            let contentJSONAsObject = asJSONObject(body: contentJSON)

            visitor.visitAnyJSONBytesField(value: contentJSONAsObject)

        }

    }

}

/// The obvious goal for Hashable/Equatable conformance would be for

/// hash and equality to behave as if we always decoded the inner

/// object and hashed or compared that.  Unfortunately, Any typically

/// stores serialized contents and we don't always have the ability to

/// deserialize it.  Since none of our supported serializations are

/// fully deterministic, we can't even ensure that equality will

/// behave this way when the Any contents are in the same

/// serialization.

///

/// As a result, we can only really perform a "best effort" equality

/// test.  Of course, regardless of the above, we must guarantee that

/// hashValue is compatible with equality.

extension AnyMessageStorage {

    // Can't use _valueData for a few reasons:

    // 1. Since decode is done on demand, two objects could be equal

    //    but created differently (one from JSON, one for Message, etc.),

    //    and the hash values have to be equal even if we don't have data

    //    yet.

    // 2. map<> serialization order is undefined. At the time of writing

    //    the Swift, Objective-C, and Go runtimes all tend to have random

    //    orders, so the messages could be identical, but in binary form

    //    they could differ.

    public func hash(into hasher: inout Hasher) {

        if !_typeURL.isEmpty {

            hasher.combine(_typeURL)

        }

    }

    func isEqualTo(other: AnyMessageStorage) -> Bool {

        if _typeURL != other._typeURL {

            return false

        }

        // Since the library does lazy Any decode, equality is a very hard problem.

        // It things exactly match, that's pretty easy, otherwise, one ends up having

        // to error on saying they aren't equal.

        //

        // The best option would be to have Message forms and compare those, as that

        // removes issues like map<> serialization order, some other protocol buffer

        // implementation details/bugs around serialized form order, etc.; but that

        // would also greatly slow down equality tests.

        //

        // Do our best to compare what is present have...

        // If both have messages, check if they are the same.

        if case .message(let myMsg) = state, case .message(let otherMsg) = other.state,

            type(of: myMsg) == type(of: otherMsg)

        {

            // Since the messages are known to be same type, we can claim both equal and

            // not equal based on the equality comparison.

            return myMsg.isEqualTo(message: otherMsg)

        }

        // If both have serialized data, and they exactly match; the messages are equal.

        // Because there could be map in the message, the fact that the data isn't the

        // same doesn't always mean the messages aren't equal. Likewise, the binary could

        // have been created by a library that doesn't order the fields, or the binary was

        // created using the appending ability in of the binary format.

        if case .binary(let myValue) = state, case .binary(let otherValue) = other.state, myValue == otherValue {

            return true

        }

        // If both have contentJSON, and they exactly match; the messages are equal.

        // Because there could be map in the message (or the JSON could just be in a different

        // order), the fact that the JSON isn't the same doesn't always mean the messages

        // aren't equal.

        if case .contentJSON(let myJSON, _) = state,

            case .contentJSON(let otherJSON, _) = other.state,

            myJSON == otherJSON

        {

            return true

        }

        // Out of options. To do more compares, the states conversions would have to be

        // done to do comparisons; and since equality can be used somewhat removed from

        // a developer (if they put protos in a Set, use them as keys to a Dictionary, etc),

        // the conversion cost might be to high for those uses.  Give up and say they aren't equal.

        return false

    }

}

// _CustomJSONCodable support for Google_Protobuf_Any

extension AnyMessageStorage {

    // Spec for Any says this should contain atleast one slash. Looking at upstream languages, most

    // actually look up the value in their runtime registries, but since we do deferred parsing

    // we can't assume the registry is complete, thus just do this minimal validation check.

    fileprivate func isTypeURLValid() -> Bool {

        _typeURL.contains(where: { $0 == "/" })

    }

    // Override the traversal-based JSON encoding

    // This builds an Any JSON representation from one of:

    //  * The message we were initialized with,

    //  * The JSON fields we last deserialized, or

    //  * The protobuf field we were deserialized from.

    // The last case requires locating the type, deserializing

    // into an object, then reserializing back to JSON.

    func encodedJSONString(options: JSONEncodingOptions) throws -> String {

        switch state {

        case .binary(let valueData):

            // Follow the C++ protostream_objectsource.cc's

            // ProtoStreamObjectSource::RenderAny() special casing of an empty value.

            if valueData.isEmpty && _typeURL.isEmpty {

                return "{}"

            }

            guard isTypeURLValid() else {

                if _typeURL.isEmpty {

                    throw SwiftProtobufError.JSONEncoding.emptyAnyTypeURL()

                }

                throw SwiftProtobufError.JSONEncoding.invalidAnyTypeURL(type_url: _typeURL)

            }

            if valueData.isEmpty {

                var jsonEncoder = JSONEncoder()

                jsonEncoder.startObject()

                jsonEncoder.startField(name: "@type")

                jsonEncoder.putStringValue(value: _typeURL)

                jsonEncoder.endObject()

                return jsonEncoder.stringResult

            }

            // Transcode by decoding the binary data to a message object

            // and then recode back into JSON.

            guard let messageType = Google_Protobuf_Any.messageType(forTypeURL: _typeURL) else {

                // If we don't have the type available, we can't decode the

                // binary value, so we're stuck.  (The Google spec does not

                // provide a way to just package the binary value for someone

                // else to decode later.)

                throw JSONEncodingError.anyTranscodeFailure

            }

            let m = try messageType.init(serializedBytes: valueData, partial: true)

            return try serializeAnyJSON(for: m, typeURL: _typeURL, options: options)

        case .message(let msg):

            // We should have been initialized with a typeURL, make sure it is valid.

            if !_typeURL.isEmpty && !isTypeURLValid() {

                throw SwiftProtobufError.JSONEncoding.invalidAnyTypeURL(type_url: _typeURL)

            }

            // If it was cleared, default it.

            let url = !_typeURL.isEmpty ? _typeURL : buildTypeURL(forMessage: msg, typePrefix: defaultAnyTypeURLPrefix)

            return try serializeAnyJSON(for: msg, typeURL: url, options: options)

        case .contentJSON(let contentJSON, _):

            guard isTypeURLValid() else {

                if _typeURL.isEmpty {

                    throw SwiftProtobufError.JSONEncoding.emptyAnyTypeURL()

                }

                throw SwiftProtobufError.JSONEncoding.invalidAnyTypeURL(type_url: _typeURL)

            }

            var jsonEncoder = JSONEncoder()

            jsonEncoder.startObject()

            jsonEncoder.startField(name: "@type")

            jsonEncoder.putStringValue(value: _typeURL)

            if !contentJSON.isEmpty {

                jsonEncoder.append(staticText: ",")

                // NOTE: This doesn't really take `options` into account since it is

                // just reflecting out what was taken in originally.

                jsonEncoder.append(utf8Bytes: contentJSON)

            }

            jsonEncoder.endObject()

            return jsonEncoder.stringResult

        }

    }

    // TODO: If the type is well-known or has already been registered,

    // we should consider decoding eagerly.  Eager decoding would

    // catch certain errors earlier (good) but would probably be

    // a performance hit if the Any contents were never accessed (bad).

    // Of course, we can't always decode eagerly (we don't always have the

    // message type available), so the deferred logic here is still needed.

    func decodeJSON(from decoder: inout JSONDecoder) throws {

        try decoder.scanner.skipRequiredObjectStart()

        // Reset state

        _typeURL = String()

        state = .binary(Data())

        if decoder.scanner.skipOptionalObjectEnd() {

            return

        }

        var jsonEncoder = JSONEncoder()

        while true {

            let key = try decoder.scanner.nextQuotedString()

            try decoder.scanner.skipRequiredColon()

            if key == "@type" {

                _typeURL = try decoder.scanner.nextQuotedString()

                guard isTypeURLValid() else {

                    throw SwiftProtobufError.JSONDecoding.invalidAnyTypeURL(type_url: _typeURL)

                }

            } else {

                jsonEncoder.startField(name: key)

                let keyValueJSON = try decoder.scanner.skip()

                jsonEncoder.append(text: keyValueJSON)

            }

            if decoder.scanner.skipOptionalObjectEnd() {

                if _typeURL.isEmpty {

                    throw SwiftProtobufError.JSONDecoding.emptyAnyTypeURL()

                }

                // Capture the options, but set the messageDepthLimit to be what

                // was left right now, as that is the limit when the JSON is finally

                // parsed.

                var updatedOptions = decoder.options

                updatedOptions.messageDepthLimit = decoder.scanner.recursionBudget

                state = .contentJSON(Array(jsonEncoder.dataResult), updatedOptions)

                return

            }

            try decoder.scanner.skipRequiredComma()

        }

    }

}