/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
 * This source file is part of the Cangjie project, licensed under Apache-2.0
 * with Runtime Library Exception.
 *
 * See https://cangjie-lang.cn/pages/LICENSE for license information.
 */

/**
 * @file
 *
 * This file implements the Fuzzer class and FuzzerBuilder class.
 */

package stdx.fuzz

import std.collection.ArrayList
import std.convert.Parsable
import std.process.*

public let FUZZ_VERSION: String = "1.0.0"
var g_fuzzer: Option<Fuzzer> = None

/**
 * To distinguish the type of target function.
 * ArrayFunc: Target function input is an UInt8 array.
 * DataProviderFunc: Target function input is `fuzz.FuzzDataProvider` object.
 */
enum FuncType {
    ArrayFunc | DataProviderFunc
}

/**
 * The core of fuzzing engine.
 */
public class Fuzzer {
    var arrayTargetFunc: (Array<UInt8>) -> Int32 = {_ => 0}
    var dpTargetFunc: (FuzzDataProvider) -> Int32 = {_ => 0}
    var ftype: FuncType
    var args: Array<String>
    var dpMaxLen: UInt32 = 4096 // default value of libfuzzer

    /**
     * For libfuzzer <= 14 ONLY.
     * For DataProvider mode ONLY.
     * When input length is not long enough, callback function will throw ExhaustedException.
     * This will cause "ERROR: no interesting inputs were found." in libfuzzer, so we should make a fake of coverage and
     * cheat libfuzzer.
     */
    var fakeCoverage: Bool = false
    var fakeCoverageArea: CPointer<UInt8> = CPointer<UInt8>()

    /**
     * For DataProvider mode ONLY.
     * If true, when FuzzDataProvider.consumeXXX is invoked, the return value will printed to stdout.
     * Default is false.
     */
    var debugDataProvider: Bool = false

    /**
     * Create a new Fuzzer.
     * Construction with targetFunction.
     * The launch arguments of fuzzer is `Process.current.arguments`.
     *
     * @param targetFunction A function input is a UInt8 Array.
     */
    public init(targetFunction: (Array<UInt8>) -> Int32) {
        this(targetFunction, Process.current.arguments)
    }

    /**
     * Create a new Fuzzer.
     * Construction with targetFunction and args.
     * The launch argument of fuzzer is @param args.
     *
     * @param targetFunction A function input is a UInt8 Array.
     * @param args The Array of args.
     */
    public init(targetFunction: (Array<UInt8>) -> Int32, args: Array<String>) {
        this.arrayTargetFunc = targetFunction
        this.args = args
        this.ftype = FuncType.ArrayFunc
    }

    /**
     * Create a new Fuzzer.
     * Construction with targetFunction.
     * The launch arguments of fuzzer is `Process.current.arguments`.
     *
     * @param targetFunction A function input is `fuzz.FuzzDataProvider` object.
     */
    public init(targetFunction: (FuzzDataProvider) -> Int32) {
        this(targetFunction, Process.current.arguments)
    }

    /**
     * Create a new Fuzzer.
     * Construction with targetFunction and args.
     * The launch argument of fuzzer is @param args.
     *
     * @param targetFunction A function input is `fuzz.FuzzDataProvider` object.
     * @param args The Array of args.
     */
    public init(targetFunction: (FuzzDataProvider) -> Int32, args: Array<String>) {
        this.dpTargetFunc = targetFunction
        this.args = args
        this.ftype = DataProviderFunc
    }

    /**
     * Create a new Fuzzer.
     * Construction with FuzzerBuilder.
     *
     * @param fb A FuzzerBuilder class.
     */
    init(fb: FuzzerBuilder) {
        this.arrayTargetFunc = fb.arrayTargetFunc
        this.dpTargetFunc = fb.dpTargetFunc
        this.args = fb.args
        this.ftype = fb.ftype
    }

    // args: getter/setter.
    public func getArgs(): Array<String> {
        return this.args
    }

    public func setArgs(args: Array<String>): Unit {
        this.args = args
    }

    // targetFunction: setter only.
    public func setTargetFunction(targetFunction: (Array<UInt8>) -> Int32): Unit {
        this.arrayTargetFunc = targetFunction
        this.ftype = ArrayFunc
    }

    public func setTargetFunction(targetFunction: (FuzzDataProvider) -> Int32): Unit {
        this.dpTargetFunc = targetFunction
        this.ftype = DataProviderFunc
    }

    /**
     * Avoid libfuzzer kill itself with "ERROR: no interesting inputs were found."
     * If enabled, an dummy coverage area will be added into libfuzzer. You should
     * make sure the fuzz target is instrumented before enable this feature.
     * See also `disableFakeCoverage`
     * For libfuzzer <= 14 ONLY.
     * For DataProvider mode ONLY.
     */
    public func enableFakeCoverage(): Unit {
        this.fakeCoverage = true
    }

    /**
     * Disable fake coverage and reset to default.
     * See also `enableFakeCoverage`
     * For libfuzzer <= 14 ONLY.
     * For DataProvider mode ONLY.
     */
    public func disableFakeCoverage(): Unit {
        this.fakeCoverage = false
    }

    /**
     * When DataProvider.consumeXXX is invoked, the return value will printed to stdout.
     * For DataProvider mode ONLY.
     */
    public func enableDebugDataProvider(): Unit {
        this.debugDataProvider = true
    }

    /**
     * When FuzzDataProvider.consumeXXX is invoked, no information will printed to stdout.
     * For DataProvider mode ONLY.
     */
    public func disableDebugDataProvider(): Unit {
        this.debugDataProvider = false
    }

    /**
     * Start fuzzing now. This function will finally call LLVMFuzzerRunDriver to run libfuzzer.
     * The fuzzing loop will stoped if an uncatched exception thrown, or running count reach limit,
     * or timeout, or OOM, and so on.
     */
    public func startFuzz(): Unit {
        // The callback function of libfuzzer cannot obtain information of "this".
        // So we must use global variable "g_fuzzer" to complete fuzzing process.
        g_fuzzer = this

        if (this.fakeCoverage) {
            this.fakeCoverageArea = LibC.malloc<UInt8>(count: 1)
            unsafe {
                __sanitizer_cov_8bit_counters_init(this.fakeCoverageArea, this.fakeCoverageArea + 1)
            }
        }

        let argc = args.size + 1
        let argv_list: ArrayList<CString> = ArrayList()

        for (arg in args) {
            let LIBFUZZER_MAX_LEN = "-max_len="
            if (arg.startsWith(LIBFUZZER_MAX_LEN)) {
                this.dpMaxLen = UInt32.parse(arg[LIBFUZZER_MAX_LEN.size..])
            }
        }
        unsafe {
            // argv[0] is a place holder
            var argv_ptr = CPointer<CPointer<UInt8>>()
            try {
                argv_list.add(LibC.mallocCString("program_name"))
                // argv[1:] is args
                for (arg in args) {
                    argv_list.add(LibC.mallocCString(arg))
                }

                // Apply for resources for libfuzzer FFI.
                var v_argc = Int32(argc)
                argv_ptr = LibC.malloc<CPointer<UInt8>>(count: argc)
                for (i in 0..argc) {
                    argv_ptr.write(i, argv_list[i].getChars())
                }

                // Call libfuzzer.
                LLVMFuzzerRunDriver(inout v_argc, inout argv_ptr, libfuzzerCallback)
            } finally {
                // Release resources of argv.
                LibC.free(argv_ptr)
                for (cstring in argv_list) {
                    LibC.free(cstring)
                }
            }
        }
    }
}

/**
 * Builder of Fuzzer.
 */
public class FuzzerBuilder {
    var arrayTargetFunc: (Array<UInt8>) -> Int32 = {_ => 0}
    var dpTargetFunc: (FuzzDataProvider) -> Int32 = {_ => 0}
    var ftype = ArrayFunc

    // By default, parameters are transferred through the command line.
    var args: Array<String> = Process.current.arguments

    /**
     * Create a new FuzzerBuilder.
     * Construction with targetFunction.
     *
     * @param targetFunction A function input is a UInt8 Array.
     */
    public init(targetFunction: (Array<UInt8>) -> Int32) {
        this.arrayTargetFunc = targetFunction
        this.ftype = ArrayFunc
    }

    /**
     * Create a new FuzzerBuilder.
     * Construction with targetFunction.
     *
     * @param targetFunction A function input is `fuzz.FuzzDataProvider` object.
     */
    public init(targetFunction: (FuzzDataProvider) -> Int32) {
        this.dpTargetFunc = targetFunction
        this.ftype = DataProviderFunc
    }

    /**
     * Reset Args.
     *
     * @param args The Array of Args.
     */
    public func setArgs(args: Array<String>): FuzzerBuilder {
        this.args = args
        return this
    }

    /**
     * Reset TargetFunction.
     *
     * @param targetFunction A function input is a UInt8 Array.
     */
    public func setTargetFunction(targetFunction: (Array<UInt8>) -> Int32): FuzzerBuilder {
        this.arrayTargetFunc = targetFunction
        this.ftype = ArrayFunc
        return this
    }

    /**
     * Reset TargetFunction.
     *
     * @param targetFunction A function input is a `fuzz.FuzzDataProvider` object.
     */
    public func setTargetFunction(targetFunction: (FuzzDataProvider) -> Int32): FuzzerBuilder {
        this.dpTargetFunc = targetFunction
        this.ftype = DataProviderFunc
        return this
    }

    /**
     * @return Fuzzer Build a new Fuzzer class.
     */
    public func build(): Fuzzer {
        return Fuzzer(this)
    }
}