// Copyright (c) 2014 - 2024 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

//

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

import Foundation

import SwiftProtobuf

public enum FuzzOption<T: SupportsFuzzOptions> {

    case boolean(WritableKeyPath<T, Bool>)

    case byte(WritableKeyPath<T, Int>, mod: UInt8 = .max)

}

public protocol SupportsFuzzOptions {

    static var fuzzOptionsList: [FuzzOption<Self>] { get }

    init()

}

extension SupportsFuzzOptions {

    public static func extractOptions(

        _ start: UnsafeRawPointer,

        _ count: Int

    ) -> (Self, UnsafeRawBufferPointer)? {

        var start = start

        let initialCount = count

        var count = count

        var options = Self()

        let reportInfo = ProcessInfo.processInfo.environment["DUMP_DECODE_INFO"] == "1"

        // No format can start with zero (invalid tag, not really UTF-8), so use that to

        // indicate there are decoding options. The one case that can start with a zero

        // would be length delimited binary, but since that's a zero length message, we

        // can go ahead and use that one also.

        guard count >= 2, start.loadUnaligned(as: UInt8.self) == 0 else {

            if reportInfo {

                print("No options to decode")

            }

            return (options, UnsafeRawBufferPointer(start: start, count: count))

        }

        // Step over the zero

        start += 1

        count -= 1

        var optionsBits = start.loadUnaligned(as: UInt8.self)

        start += 1

        count -= 1

        var bit = 0

        for opt in fuzzOptionsList {

            var isSet = optionsBits & (1 << bit) != 0

            if bit == 7 {

                // About the use the last bit of this byte, to allow more options in

                // the future, use this bit to indicate reading another byte.

                guard isSet else {

                    // No continuation, just return whatever we got.

                    bit = 8

                    break

                }

                guard count >= 1 else {

                    return nil  // No data left to read bits

                }

                optionsBits = start.loadUnaligned(as: UInt8.self)

                start += 1

                count -= 1

                bit = 0

                isSet = optionsBits & (1 << bit) != 0

            }

            switch opt {

            case .boolean(let keypath):

                options[keyPath: keypath] = isSet

            case .byte(let keypath, let mod):

                assert(mod >= 1 && mod <= UInt8.max)

                if isSet {

                    guard count >= 1 else {

                        return nil  // No more bytes to get a value, fail

                    }

                    let value = start.loadUnaligned(as: UInt8.self)

                    start += 1

                    count -= 1

                    options[keyPath: keypath] = Int(value % mod)

                }

            }

            bit += 1

        }

        // Ensure the any remaining bits are zero so they can be used in the future

        while bit < 8 {

            if optionsBits & (1 << bit) != 0 {

                return nil

            }

            bit += 1

        }

        if reportInfo {

            print("\(initialCount - count) bytes consumed off front for options: \(options)")

        }

        return (options, UnsafeRawBufferPointer(start: start, count: count))

    }

}

extension BinaryDecodingOptions: SupportsFuzzOptions {

    public static var fuzzOptionsList: [FuzzOption<Self>] {

        [

            // NOTE: Do not reorder these in the future as it invalidates all

            // existing cases.

            // The default depth is 100, so limit outselves to modding by 8 to

            // avoid allowing larger depths that could timeout.

            .byte(\.messageDepthLimit, mod: 8),

            .boolean(\.discardUnknownFields),

        ]

    }

}

extension JSONDecodingOptions: SupportsFuzzOptions {

    public static var fuzzOptionsList: [FuzzOption<Self>] {

        [

            // NOTE: Do not reorder these in the future as it invalidates all

            // existing cases.

            // The default depth is 100, so limit outselves to modding by 8 to

            // avoid allowing larger depths that could timeout.

            .byte(\.messageDepthLimit, mod: 8),

            .boolean(\.ignoreUnknownFields),

        ]

    }

}

extension TextFormatDecodingOptions: SupportsFuzzOptions {

    public static var fuzzOptionsList: [FuzzOption<Self>] {

        [

            // NOTE: Do not reorder these in the future as it invalidates all

            // existing cases.

            // The default depth is 100, so limit outselves to modding by 8 to

            // avoid allowing larger depths that could timeout.

            .byte(\.messageDepthLimit, mod: 8),

            .boolean(\.ignoreUnknownFields),

            .boolean(\.ignoreUnknownExtensionFields),

        ]

    }

}