#include <iostream>
#include <vector>
#include <getopt.h>
#include <unistd.h>
#include <dirent.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <stdio.h>
#include <stdlib.h>
#include <time.h>
#include <sys/wait.h>
#include <dlfcn.h>
#include <cstdio>
#include <string>
#include <iomanip>
#include <string.h>
#include <list>
#include <algorithm>
#include <stdarg.h>
#include "acl/acl.h"
#include "aclnn/aclnn_base.h"
#include "acl/acl_base.h"
#include "acl/acl_rt.h"
#include "sec_common.h"
#include "utils.h"

#include "../fuzzlib/secodeFuzz.h"
#include "aclnnop/aclnn_bucketize.h"

using namespace std;

static std::string ShapeToStr(const std::vector<int64_t>& s) {
    std::string r = "[";
    for (size_t k = 0; k < s.size(); k++) {
        if (k) r += ",";
        r += std::to_string(s[k]);
    }
    return r + "]";
}
int aclnnBucketize_Normal(int fuzz_num)
{
    auto ret = Init(deviceId, &context, &stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);
    WriteDTLog("aclnnBucketize", "aclnnBucketize op fuzz start \n");
    const long long kSingleInputSizeMax = 1073741824LL;
    const long long kTotalInputSizeMax = 2147483648LL;

    // dtype enums: FLOAT=0 FLOAT16=1 BF16=27 INT8=2 INT16=6 INT32=3 INT64=9 UINT8=4
    int dtype0Enum[] = { 0, 1, 27, 2, 6, 3, 9, 4 };
    int dtype0Num = sizeof(dtype0Enum) / sizeof(dtype0Enum[0]);

    aclrtSetOpExecuteTimeOut(1091);
    DT_Set_TimeOut_Second(3000);
    DT_FUZZ_START(0, fuzz_num, "aclnnBucketize", 0)
    {
        int i = 0;
        int shapeDimMax = 16;
        int aclnnRet = -1;  // aclnn 实际调用返回值（成功判定依据）
        void* selfDeviceAddr = nullptr;
        aclTensor* self = nullptr;
        void* boundariesDeviceAddr = nullptr;
        aclTensor* boundaries = nullptr;
        void* outDeviceAddr = nullptr;
        aclTensor* out = nullptr;
        uint64_t workspaceSize = 0;
        aclOpExecutor* executor;
        void* workspaceAddr = nullptr;

        LOG_PRINT("aclnnBucketize_Normal %d\n", fuzzi);
        printf("[FUZZ_LOG] plog_file: aclnnBucketize_fuzz_plog_case_%d.log\n", fuzzi);
        aclnnRet = -1;
        try
        {
        // DT mutation data
            int sizeBucket = *(int *)DTSetGetNumberRange(&g_Element[i++], 1, 0, 3, "i");
            int dim_num0 = *(int *)DTSetGetNumberRange(&g_Element[i++], 2, 0, 8, "i");
            {
                const long long sizeTargets[4] = {1000LL, 65536LL, 262144LL, 524288LL};
                int nTensors = 2;
                long long target = sizeTargets[sizeBucket] / nTensors;
                if (target < 1) { target = 1; }
                int nd = (dim_num0 > 0) ? dim_num0 : 1;
                long long perDim = (long long)pow((double)target, 1.0 / nd);
                if (perDim < 1) { perDim = 1; }
                if (perDim > MAX_SHAPE_SIZE) { perDim = MAX_SHAPE_SIZE; }
                shapeDimMax = (int)perDim;
            }
            int shape1 = *(int *)DTSetGetNumberRange(&g_Element[i++], 2, 1, shapeDimMax, "i");
            int shape2 = *(int *)DTSetGetNumberRange(&g_Element[i++], 2, 1, shapeDimMax, "i");
            int shape3 = *(int *)DTSetGetNumberRange(&g_Element[i++], 2, 1, shapeDimMax, "i");
            int shape4 = *(int *)DTSetGetNumberRange(&g_Element[i++], 2, 1, shapeDimMax, "i");
            int shape5 = *(int *)DTSetGetNumberRange(&g_Element[i++], 2, 1, shapeDimMax, "i");
            int shape6 = *(int *)DTSetGetNumberRange(&g_Element[i++], 2, 1, shapeDimMax, "i");
            int shape7 = *(int *)DTSetGetNumberRange(&g_Element[i++], 2, 1, shapeDimMax, "i");
            int shape8 = *(int *)DTSetGetNumberRange(&g_Element[i++], 2, 1, shapeDimMax, "i");
            auto dtype0 = (aclDataType)*(int *)DT_SetGetNumberEnumV3(i++, aclDataType::ACL_FLOAT, dtype0Enum, dtype0Num);
            int boundariesSize = *(int *)DTSetGetNumberRange(&g_Element[i++], 2, 1, 100, "i");
            bool outInt32 = (bool)*(int *)DTSetGetNumberRange(&g_Element[i++], 0, 0, 1, "i");
            bool right = (bool)*(int *)DTSetGetNumberRange(&g_Element[i++], 0, 0, 1, "i");
            std::vector<int64_t> selfShape = {};
            if (dim_num0 == 1) {
                selfShape = {shape1};
            } else if (dim_num0 == 2) {
                selfShape = {shape1, shape2};
            } else if (dim_num0 == 3) {
                selfShape = {shape1, shape2, shape3};
            } else if (dim_num0 == 4) {
                selfShape = {shape1, shape2, shape3, shape4};
            } else if (dim_num0 == 5) {
                selfShape = {shape1, shape2, shape3, shape4, shape5};
            } else if (dim_num0 == 6) {
                selfShape = {shape1, shape2, shape3, shape4, shape5, shape6};
            } else if (dim_num0 == 7) {
                selfShape = {shape1, shape2, shape3, shape4, shape5, shape6, shape7};
            } else if (dim_num0 == 8) {
                selfShape = {shape1, shape2, shape3, shape4, shape5, shape6, shape7, shape8};
            }
            std::vector<int64_t> boundariesShape = { boundariesSize };
            std::vector<int64_t> outShape = selfShape;
            long long selfSize = (long long)GetShapeSize(selfShape);
            long long outSize = (long long)GetShapeSize(outShape);
            long long shape_size = (selfSize > outSize) ? selfSize : outSize;
            if (shape_size > kSingleInputSizeMax) { printf("shape_size too big\n"); goto ERROR; }
            if (dtype0 == aclDataType::ACL_FLOAT) {
                std::vector<float> selfHostData(selfSize, 1);
                ret = CreateAclTensor(selfHostData, selfShape, &selfDeviceAddr, dtype0, &self);
            CHECK_RET(ret == ACL_SUCCESS, goto ERROR);
            } else if (dtype0 == aclDataType::ACL_FLOAT16) {
                std::vector<int16_t> selfHostData(selfSize, 1);
                ret = CreateAclTensor(selfHostData, selfShape, &selfDeviceAddr, dtype0, &self);
            CHECK_RET(ret == ACL_SUCCESS, goto ERROR);
            } else if (dtype0 == aclDataType::ACL_BF16) {
                std::vector<int16_t> selfHostData(selfSize, 1);
                ret = CreateAclTensor(selfHostData, selfShape, &selfDeviceAddr, dtype0, &self);
            CHECK_RET(ret == ACL_SUCCESS, goto ERROR);
            } else if (dtype0 == aclDataType::ACL_INT8) {
                std::vector<int8_t> selfHostData(selfSize, 1);
                ret = CreateAclTensor(selfHostData, selfShape, &selfDeviceAddr, dtype0, &self);
            CHECK_RET(ret == ACL_SUCCESS, goto ERROR);
            } else if (dtype0 == aclDataType::ACL_INT16) {
                std::vector<int16_t> selfHostData(selfSize, 1);
                ret = CreateAclTensor(selfHostData, selfShape, &selfDeviceAddr, dtype0, &self);
            CHECK_RET(ret == ACL_SUCCESS, goto ERROR);
            } else if (dtype0 == aclDataType::ACL_INT32) {
                std::vector<int32_t> selfHostData(selfSize, 1);
                ret = CreateAclTensor(selfHostData, selfShape, &selfDeviceAddr, dtype0, &self);
            CHECK_RET(ret == ACL_SUCCESS, goto ERROR);
            } else if (dtype0 == aclDataType::ACL_INT64) {
                std::vector<int64_t> selfHostData(selfSize, 1);
                ret = CreateAclTensor(selfHostData, selfShape, &selfDeviceAddr, dtype0, &self);
            CHECK_RET(ret == ACL_SUCCESS, goto ERROR);
            } else if (dtype0 == aclDataType::ACL_UINT8) {
                std::vector<uint8_t> selfHostData(selfSize, 1);
                ret = CreateAclTensor(selfHostData, selfShape, &selfDeviceAddr, dtype0, &self);
            CHECK_RET(ret == ACL_SUCCESS, goto ERROR);
            }
            if (dtype0 == aclDataType::ACL_FLOAT) {
                std::vector<float> boundariesHostData(boundariesSize, 1);
                for (int64_t k = 0; k < boundariesSize; k++) { boundariesHostData[k] = (float)(k + 1); }
                ret = CreateAclTensor(boundariesHostData, boundariesShape, &boundariesDeviceAddr, dtype0, &boundaries);
            CHECK_RET(ret == ACL_SUCCESS, goto ERROR);
            } else if (dtype0 == aclDataType::ACL_FLOAT16) {
                std::vector<int16_t> boundariesHostData(boundariesSize, 1);
                for (int64_t k = 0; k < boundariesSize; k++) { boundariesHostData[k] = (int16_t)(k + 1); }
                ret = CreateAclTensor(boundariesHostData, boundariesShape, &boundariesDeviceAddr, dtype0, &boundaries);
            CHECK_RET(ret == ACL_SUCCESS, goto ERROR);
            } else if (dtype0 == aclDataType::ACL_BF16) {
                std::vector<int16_t> boundariesHostData(boundariesSize, 1);
                for (int64_t k = 0; k < boundariesSize; k++) { boundariesHostData[k] = (int16_t)(k + 1); }
                ret = CreateAclTensor(boundariesHostData, boundariesShape, &boundariesDeviceAddr, dtype0, &boundaries);
            CHECK_RET(ret == ACL_SUCCESS, goto ERROR);
            } else if (dtype0 == aclDataType::ACL_INT8) {
                std::vector<int8_t> boundariesHostData(boundariesSize, 1);
                for (int64_t k = 0; k < boundariesSize; k++) { boundariesHostData[k] = (int8_t)(k + 1); }
                ret = CreateAclTensor(boundariesHostData, boundariesShape, &boundariesDeviceAddr, dtype0, &boundaries);
            CHECK_RET(ret == ACL_SUCCESS, goto ERROR);
            } else if (dtype0 == aclDataType::ACL_INT16) {
                std::vector<int16_t> boundariesHostData(boundariesSize, 1);
                for (int64_t k = 0; k < boundariesSize; k++) { boundariesHostData[k] = (int16_t)(k + 1); }
                ret = CreateAclTensor(boundariesHostData, boundariesShape, &boundariesDeviceAddr, dtype0, &boundaries);
            CHECK_RET(ret == ACL_SUCCESS, goto ERROR);
            } else if (dtype0 == aclDataType::ACL_INT32) {
                std::vector<int32_t> boundariesHostData(boundariesSize, 1);
                for (int64_t k = 0; k < boundariesSize; k++) { boundariesHostData[k] = (int32_t)(k + 1); }
                ret = CreateAclTensor(boundariesHostData, boundariesShape, &boundariesDeviceAddr, dtype0, &boundaries);
            CHECK_RET(ret == ACL_SUCCESS, goto ERROR);
            } else if (dtype0 == aclDataType::ACL_INT64) {
                std::vector<int64_t> boundariesHostData(boundariesSize, 1);
                for (int64_t k = 0; k < boundariesSize; k++) { boundariesHostData[k] = (int64_t)(k + 1); }
                ret = CreateAclTensor(boundariesHostData, boundariesShape, &boundariesDeviceAddr, dtype0, &boundaries);
            CHECK_RET(ret == ACL_SUCCESS, goto ERROR);
            } else if (dtype0 == aclDataType::ACL_UINT8) {
                std::vector<uint8_t> boundariesHostData(boundariesSize, 1);
                for (int64_t k = 0; k < boundariesSize; k++) { boundariesHostData[k] = (uint8_t)(k + 1); }
                ret = CreateAclTensor(boundariesHostData, boundariesShape, &boundariesDeviceAddr, dtype0, &boundaries);
            CHECK_RET(ret == ACL_SUCCESS, goto ERROR);
            }
            if (outInt32) {
                std::vector<int32_t> outHostData(outSize, 1);
                ret = CreateAclTensor(outHostData, outShape, &outDeviceAddr, aclDataType::ACL_INT32, &out);
            CHECK_RET(ret == ACL_SUCCESS, goto ERROR);
            } else {
                std::vector<int64_t> outHostData(outSize, 1);
                ret = CreateAclTensor(outHostData, outShape, &outDeviceAddr, aclDataType::ACL_INT64, &out);
            CHECK_RET(ret == ACL_SUCCESS, goto ERROR);
            }

            ret = aclnnBucketizeGetWorkspaceSize(self, boundaries, outInt32, right, out, &workspaceSize, &executor);
printf("[FUZZ_PARAM] self: dtype=%d, shape=%s; boundaries: dtype=%d, shape=%s; out: dtype=%d, shape=%s; outInt32=%d; right=%d\n", (int)dtype0, ShapeToStr(selfShape).c_str(), (int)dtype0, ShapeToStr(boundariesShape).c_str(), (int)(outInt32 ? aclDataType::ACL_INT32 : aclDataType::ACL_INT64), ShapeToStr(outShape).c_str(), (int)outInt32, (int)right);
            CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnBucketizeGetWorkspaceSize failed. ERROR: %d\n", ret); goto ERROR);
            if (workspaceSize > 0) {
                ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
                CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret); goto ERROR);
            }
            ret = aclnnBucketize(workspaceAddr, workspaceSize, executor, stream);
            aclnnRet = (int)ret;
            if (ret == ACL_SUCCESS) { printf("aclnnBucketize execute success :%d\n", (int)ret); }
            else { printf("aclnnBucketize execute failed :%d\n", (int)ret); goto ERROR; }
            ret = aclrtSynchronizeStream(stream);
            CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); goto ERROR);
            LOG_PRINT("aclnnBucketize execute success : %d\n", ret);
        }
        catch (std::exception& e)
        {
            printf("catch exception %s\n", e.what());
        }

ERROR:
        ret = aclrtSynchronizeStream(stream);
        if (aclnnRet == -1) { printf("[FUZZ_RESULT] case_%d: skip, ret=-1\n", fuzzi); }
        else { printf("[FUZZ_RESULT] case_%d: %s, ret=%d\n", fuzzi, (aclnnRet == 0) ? "success" : "faild", aclnnRet); }
        printf("end \n");
        CheckAndLog(ret);
        pid_t pid = getpid();
        std::string cmd = "grep -rinE 'ai(vec|core) error' /root/ascend/log/debug/plog/ ";
        printf("system return value %d\n", CheckResult(pid));
        if (self) { aclDestroyTensor(self); self = nullptr; }
        if (boundaries) { aclDestroyTensor(boundaries); boundaries = nullptr; }
        if (out) { aclDestroyTensor(out); out = nullptr; }
        if (selfDeviceAddr) { aclrtFree(selfDeviceAddr); selfDeviceAddr = nullptr; }
        if (boundariesDeviceAddr) { aclrtFree(boundariesDeviceAddr); boundariesDeviceAddr = nullptr; }
        if (outDeviceAddr) { aclrtFree(outDeviceAddr); outDeviceAddr = nullptr; }
        if (workspaceSize > 0 && workspaceAddr) { aclrtFree(workspaceAddr); workspaceAddr = nullptr; }
    }
    DT_FUZZ_END()
    WriteDTLog("aclnnBucketize", "aclnnBucketize op fuzz total number:%u, para error number %u, execute error number %u \n", fuzz_num, aclnn_INVALID_Para, aclnn_Inner_Error + Runtime_Error);
    WriteDTLog("aclnnBucketize", "aclnnBucketize op fuzz end \n");
    aclrtDestroyStream(stream);
    aclrtDestroyContext(context);
    aclrtResetDevice(deviceId);
    aclFinalize();
    return 0;
}

extern "C" int aclnnBucketize_fuzz(int argc, char *argv[]);
int aclnnBucketize_fuzz(int argc, char *argv[])
{
    int fuzz_num = MAXFUZZNUMBER;
    if (argc < 1) {
        cout << "[ERROR] input arg num must be bigger than 1 \n" << endl;
        return -1;
    } else if (argc > 1) {
        fuzz_num = atoi(argv[2]);
    }
    if (argc == 4) { deviceId = atoi(argv[3]); }
    printf("fuzz_num %d \n", fuzz_num);
    printf("aclnnBucketize_Normal :%d \n", aclnnBucketize_Normal(fuzz_num));
    return 0;
}