#include <algorithm>
#include <vector>
#include <string>
#include <cstring>
#include <iostream>
#include <unistd.h>
#include <sys/wait.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_add_layer_norm_quant_v2.h"
using namespace std;

static int kXFamilyEnum[] = {0, 1, 27};      // FLOAT, FLOAT16, BF16
static int kComboEnum[] = {0, 1, 2, 3, 4, 5, 6};  // 0-3 static / 4-6 dynamic
static int kPairSelEnum[] = {0, 1, 2};
static int kEpsEnum[] = {0, 1, 2};

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 aclnnAddLayerNormQuantV2_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("aclnnAddLayerNormQuantV2", "aclnnAddLayerNormQuantV2 op fuzz start \n");
    const long long kSingleInputSizeMax = 1048576LL;

    aclrtSetOpExecuteTimeOut(1091);
    DT_Set_TimeOut_Second(3000);
    DT_FUZZ_START(0, fuzz_num, "aclnnAddLayerNormQuantV2", 0)
    {
        int i = 0;
        int aclnnRet = -1;
        void* x1Addr = nullptr; aclTensor* x1 = nullptr;
        void* x2Addr = nullptr; aclTensor* x2 = nullptr;
        void* gammaAddr = nullptr; aclTensor* gamma = nullptr;
        void* betaAddr = nullptr; aclTensor* beta = nullptr;
        void* biasAddr = nullptr; aclTensor* biasOptional = nullptr;
        void* s1Addr = nullptr; aclTensor* scales1Optional = nullptr;
        void* s2Addr = nullptr; aclTensor* scales2Optional = nullptr;
        void* z1Addr = nullptr; aclTensor* zeroPoints1Optional = nullptr;
        void* z2Addr = nullptr; aclTensor* zeroPoints2Optional = nullptr;
        void* y1Addr = nullptr; aclTensor* y1Out = nullptr;
        void* y2Addr = nullptr; aclTensor* y2Out = nullptr;
        void* xOutAddr = nullptr; aclTensor* xOut = nullptr;
        void* lnResAddr = nullptr; aclTensor* layernormRes = nullptr;
        void* os1Addr = nullptr; aclTensor* outScales1Out = nullptr;
        void* os2Addr = nullptr; aclTensor* outScales2Out = nullptr;
        uint64_t workspaceSize = 0;
        aclOpExecutor* executor = nullptr;
        void* workspaceAddr = nullptr;

        LOG_PRINT("aclnnAddLayerNormQuantV2_Normal %d\n", fuzzi);
        printf("[FUZZ_LOG] plog_file: aclnnAddLayerNormQuantV2_fuzz_plog_case_%d.log\n", fuzzi);
        aclnnRet = -1;
        try
        {
            // size distribution: 0=small 1=mid 2=large 3=xlarge
            int sizeBucket = *(int *)DTSetGetNumberRange(&g_Element[i++], 1, 0, 3, "i");
            int dimNum = *(int *)DTSetGetNumberRange(&g_Element[i++], 2, 2, 8, "i");      // x1 维度 2~8（dynamic 要求≥2）
            int gDimLen = *(int *)DTSetGetNumberRange(&g_Element[i++], 2, 1, 8, "i");     // gamma 后缀维度数
            int biasMode = *(int *)DTSetGetNumberRange(&g_Element[i++], 1, 0, 2, "i");    // 0=无 1=gamma形 2=x1形
            int comboIdx = *(int *)DT_SetGetNumberEnumV3(i++, 0, kComboEnum, 7);
            int addOut = *(int *)DTSetGetNumberRange(&g_Element[i++], 1, 0, 1, "i");
            int divMode = *(int *)DTSetGetNumberRange(&g_Element[i++], 1, 0, 1, "i");
            const long long sizeTargets[4] = {1000LL, 65536LL, 262144LL, 1048576LL};
            long long target = sizeTargets[sizeBucket];
            long long perDim = (long long)pow((double)target, 1.0 / dimNum);
            if (perDim < 2) { perDim = 2; }
            if (perDim > 1024) { perDim = 1024; }
            int d1 = *(int *)DTSetGetNumberRange(&g_Element[i++], 2, 1, (int)perDim, "i");
            int d2 = *(int *)DTSetGetNumberRange(&g_Element[i++], 2, 1, (int)perDim, "i");
            int d3 = *(int *)DTSetGetNumberRange(&g_Element[i++], 2, 1, (int)perDim, "i");
            int d4 = *(int *)DTSetGetNumberRange(&g_Element[i++], 2, 1, (int)perDim, "i");
            int d5 = *(int *)DTSetGetNumberRange(&g_Element[i++], 2, 1, (int)perDim, "i");
            int d6 = *(int *)DTSetGetNumberRange(&g_Element[i++], 2, 1, (int)perDim, "i");
            int d7 = *(int *)DTSetGetNumberRange(&g_Element[i++], 2, 1, (int)perDim, "i");
            int d8 = *(int *)DTSetGetNumberRange(&g_Element[i++], 2, 1, (int)perDim, "i");
            // x family dtype: static {F32,F16,BF16} / dynamic {F16,BF16}
            bool isDynamic = (comboIdx >= 4);
            int xVal;
            if (isDynamic) {
                int dynEnum[] = {0, 1};  // 枚举表值占位，取值域 {1,27}
                int dv = *(int *)DT_SetGetNumberEnumV3(i++, 0, kPairSelEnum, 2);
                xVal = (dv == 0) ? 1 : 27;
            } else {
                xVal = *(int *)DT_SetGetNumberEnumV3(i++, 0, kXFamilyEnum, 3);
            }
            aclDataType dtypeX = (aclDataType)xVal;
            // scale family dtype: 0=同x 1=F32（仅 static 且 x≠F32 时 F32 可选）
            int scaleSel = *(int *)DT_SetGetNumberEnumV3(i++, 0, kPairSelEnum, 2);
            aclDataType dtypeS = dtypeX;
            if (!isDynamic && xVal != 0 && scaleSel == 1) { dtypeS = ACL_FLOAT; }
            (void)kEpsEnum;
            int epsSel = *(int *)DTSetGetNumberRange(&g_Element[i++], 1, 0, 2, "i");
            double epsilon = (epsSel == 0) ? 1e-6 : ((epsSel == 1) ? 1e-3 : 0.1);
            const char* quantMode = isDynamic ? "dynamic" : "static";
            bool additionalOutput = (addOut == 1);
            bool divModeB = (divMode == 1);

            // 可选组合（合法表）
            bool hasS1, hasS2, hasZ1, hasZ2;
            switch (comboIdx) {
                case 0: hasS1 = true;  hasS2 = true;  hasZ1 = true;  hasZ2 = true;  break;
                case 1: hasS1 = true;  hasS2 = true;  hasZ1 = false; hasZ2 = false; break;
                case 2: hasS1 = true;  hasS2 = false; hasZ1 = true;  hasZ2 = false; break;
                case 3: hasS1 = true;  hasS2 = false; hasZ1 = false; hasZ2 = false; break;
                case 4: hasS1 = true;  hasS2 = true;  hasZ1 = false; hasZ2 = false; break;
                case 5: hasS1 = true;  hasS2 = false; hasZ1 = false; hasZ2 = false; break;
                default: hasS1 = false; hasS2 = false; hasZ1 = false; hasZ2 = false; break;
            }

            // shapes
            int dimsArr[8] = {d1, d2, d3, d4, d5, d6, d7, d8};
            std::vector<int64_t> x1Shape;
            for (int k = 0; k < dimNum; k++) { if (dimsArr[k] < 1) dimsArr[k] = 1; x1Shape.push_back(dimsArr[k]); }
            if (isDynamic && x1Shape[0] > 1) { x1Shape[0] = 1; }  // A5 bug已记录: dynamic首维>1 tiling失败
            int gLen = gDimLen;
            if (gLen < 1) gLen = 1;
            if (gLen > dimNum) gLen = dimNum;
            if (isDynamic && gLen == 1) { gLen = 2; }  // dynamic: gamma 维度≠1
            std::vector<int64_t> gammaShape;
            for (int k = dimNum - gLen; k < dimNum; k++) { gammaShape.push_back(x1Shape[k]); }
            std::vector<int64_t> betaShape = gammaShape;
            std::vector<int64_t> biasShape = (biasMode == 2) ? x1Shape : gammaShape;
            std::vector<int64_t> scalesShape = gammaShape;  // A5 perTensor[1]路径损坏(探针+样例证实)，固定gamma形
            std::vector<int64_t> xOutShape = x1Shape;
            std::vector<int64_t> lnResShape = x1Shape;
            std::vector<int64_t> yShape = x1Shape;
            std::vector<int64_t> osShape;  // x1 去掉最后一维
            for (int k = 0; k < dimNum - 1; k++) { osShape.push_back(x1Shape[k]); }

            long long x1Size = GetShapeSize(x1Shape);
            if (x1Size > kSingleInputSizeMax) { printf("shape_size too big\n"); goto ERROR; }
            long long gammaSize = GetShapeSize(gammaShape);
            long long biasSize = GetShapeSize(biasShape);
            long long scalesSize = GetShapeSize(scalesShape);
            long long osSize = GetShapeSize(osShape);

            // create tensors
            {
                std::vector<float> xh((size_t)x1Size, 1.5f);
                ret = CreateAclTensor_New(xh, x1Shape, &x1Addr, dtypeX, aclFormat::ACL_FORMAT_ND, &x1);
                CHECK_RET(ret == ACL_SUCCESS, goto ERROR);
                ret = CreateAclTensor_New(xh, x1Shape, &x2Addr, dtypeX, aclFormat::ACL_FORMAT_ND, &x2);
                CHECK_RET(ret == ACL_SUCCESS, goto ERROR);
                std::vector<float> gh((size_t)gammaSize, 1.0f);
                ret = CreateAclTensor_New(gh, gammaShape, &gammaAddr, dtypeX, aclFormat::ACL_FORMAT_ND, &gamma);
                CHECK_RET(ret == ACL_SUCCESS, goto ERROR);
                ret = CreateAclTensor_New(gh, betaShape, &betaAddr, dtypeX, aclFormat::ACL_FORMAT_ND, &beta);
                CHECK_RET(ret == ACL_SUCCESS, goto ERROR);
                if (biasMode != 0) {
                    std::vector<float> bh((size_t)biasSize, 0.5f);
                    ret = CreateAclTensor_New(bh, biasShape, &biasAddr, dtypeX, aclFormat::ACL_FORMAT_ND, &biasOptional);
                    CHECK_RET(ret == ACL_SUCCESS, goto ERROR);
                }
                if (hasS1) {
                    std::vector<float> sh((size_t)scalesSize, 0.01f);
                    ret = CreateAclTensor_New(sh, scalesShape, &s1Addr, dtypeS, aclFormat::ACL_FORMAT_ND, &scales1Optional);
                    CHECK_RET(ret == ACL_SUCCESS, goto ERROR);
                }
                if (hasS2) {
                    std::vector<float> sh((size_t)scalesSize, 0.01f);
                    ret = CreateAclTensor_New(sh, scalesShape, &s2Addr, dtypeS, aclFormat::ACL_FORMAT_ND, &scales2Optional);
                    CHECK_RET(ret == ACL_SUCCESS, goto ERROR);
                }
                if (hasZ1) {
                    std::vector<float> zh((size_t)scalesSize, 0.0f);
                    ret = CreateAclTensor_New(zh, scalesShape, &z1Addr, dtypeS, aclFormat::ACL_FORMAT_ND, &zeroPoints1Optional);
                    CHECK_RET(ret == ACL_SUCCESS, goto ERROR);
                }
                if (hasZ2) {
                    std::vector<float> zh((size_t)scalesSize, 0.0f);
                    ret = CreateAclTensor_New(zh, scalesShape, &z2Addr, dtypeS, aclFormat::ACL_FORMAT_ND, &zeroPoints2Optional);
                    CHECK_RET(ret == ACL_SUCCESS, goto ERROR);
                }
                std::vector<float> yh((size_t)x1Size, 0.0f);
                ret = CreateAclTensor_New(yh, yShape, &y1Addr, ACL_INT8, aclFormat::ACL_FORMAT_ND, &y1Out);
                CHECK_RET(ret == ACL_SUCCESS, goto ERROR);
                ret = CreateAclTensor_New(yh, yShape, &y2Addr, ACL_INT8, aclFormat::ACL_FORMAT_ND, &y2Out);
                CHECK_RET(ret == ACL_SUCCESS, goto ERROR);
                ret = CreateAclTensor_New(yh, xOutShape, &xOutAddr, dtypeX, aclFormat::ACL_FORMAT_ND, &xOut);
                CHECK_RET(ret == ACL_SUCCESS, goto ERROR);
                ret = CreateAclTensor_New(yh, lnResShape, &lnResAddr, dtypeX, aclFormat::ACL_FORMAT_ND, &layernormRes);
                CHECK_RET(ret == ACL_SUCCESS, goto ERROR);
                std::vector<float> oh((size_t)osSize, 0.0f);
                ret = CreateAclTensor_New(oh, osShape, &os1Addr, ACL_FLOAT, aclFormat::ACL_FORMAT_ND, &outScales1Out);
                CHECK_RET(ret == ACL_SUCCESS, goto ERROR);
                ret = CreateAclTensor_New(oh, osShape, &os2Addr, ACL_FLOAT, aclFormat::ACL_FORMAT_ND, &outScales2Out);
                CHECK_RET(ret == ACL_SUCCESS, goto ERROR);
            }
            ret = aclnnAddLayerNormQuantV2GetWorkspaceSize(x1, x2, gamma, beta, biasOptional,
                scales1Optional, scales2Optional, zeroPoints1Optional, zeroPoints2Optional,
                quantMode, epsilon, additionalOutput, divModeB,
                y1Out, y2Out, xOut, layernormRes, outScales1Out, outScales2Out,
                &workspaceSize, &executor);
            printf("[FUZZ_PARAM] quantMode=%s, additionalOutput=%d, divMode=%d, biasMode=%d, combo=%d, dtypeX=%d, dtypeS=%d, eps=%g; x1: dtype=%d, shape=%s; gamma: shape=%s; scales: shape=%s; y1: shape=%s; outScales: shape=%s\n",
                   quantMode, (int)additionalOutput, (int)divModeB, biasMode, comboIdx, (int)dtypeX, (int)dtypeS, epsilon,
                   (int)dtypeX, ShapeToStr(x1Shape).c_str(), ShapeToStr(gammaShape).c_str(), ShapeToStr(scalesShape).c_str(),
                   ShapeToStr(yShape).c_str(), ShapeToStr(osShape).c_str());
            CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnAddLayerNormQuantV2GetWorkspaceSize 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 = aclnnAddLayerNormQuantV2(workspaceAddr, workspaceSize, executor, stream);
            aclnnRet = (int)ret;
            if (ret == ACL_SUCCESS) { printf("aclnnAddLayerNormQuantV2 execute success :%d\n", (int)ret); }
            else { printf("aclnnAddLayerNormQuantV2 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("aclnnAddLayerNormQuantV2 execute success : %d\n", ret);
        }
        catch (std::exception& e)
        {
            printf("catch exception %s\n", e.what());
        }

ERROR:
        ret = aclrtSynchronizeStream(stream);
        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 (x1) { aclDestroyTensor(x1); }
        if (x2) { aclDestroyTensor(x2); }
        if (gamma) { aclDestroyTensor(gamma); }
        if (beta) { aclDestroyTensor(beta); }
        if (biasOptional) { aclDestroyTensor(biasOptional); }
        if (scales1Optional) { aclDestroyTensor(scales1Optional); }
        if (scales2Optional) { aclDestroyTensor(scales2Optional); }
        if (zeroPoints1Optional) { aclDestroyTensor(zeroPoints1Optional); }
        if (zeroPoints2Optional) { aclDestroyTensor(zeroPoints2Optional); }
        if (y1Out) { aclDestroyTensor(y1Out); }
        if (y2Out) { aclDestroyTensor(y2Out); }
        if (xOut) { aclDestroyTensor(xOut); }
        if (layernormRes) { aclDestroyTensor(layernormRes); }
        if (outScales1Out) { aclDestroyTensor(outScales1Out); }
        if (outScales2Out) { aclDestroyTensor(outScales2Out); }
        if (x1Addr) { aclrtFree(x1Addr); }
        if (x2Addr) { aclrtFree(x2Addr); }
        if (gammaAddr) { aclrtFree(gammaAddr); }
        if (betaAddr) { aclrtFree(betaAddr); }
        if (biasAddr) { aclrtFree(biasAddr); }
        if (s1Addr) { aclrtFree(s1Addr); }
        if (s2Addr) { aclrtFree(s2Addr); }
        if (z1Addr) { aclrtFree(z1Addr); }
        if (z2Addr) { aclrtFree(z2Addr); }
        if (y1Addr) { aclrtFree(y1Addr); }
        if (y2Addr) { aclrtFree(y2Addr); }
        if (xOutAddr) { aclrtFree(xOutAddr); }
        if (lnResAddr) { aclrtFree(lnResAddr); }
        if (os1Addr) { aclrtFree(os1Addr); }
        if (os2Addr) { aclrtFree(os2Addr); }
        if (workspaceSize > 0 && workspaceAddr) { aclrtFree(workspaceAddr); }
    }
    DT_FUZZ_END()
    WriteDTLog("aclnnAddLayerNormQuantV2", "aclnnAddLayerNormQuantV2 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("aclnnAddLayerNormQuantV2", "aclnnAddLayerNormQuantV2 op fuzz end \n");
    aclrtDestroyStream(stream);
    aclrtDestroyContext(context);
    aclrtResetDevice(deviceId);
    aclFinalize();
    return 0;
}

extern "C" int aclnnAddLayerNormQuantV2_fuzz(int argc, char *argv[])
{
    int fuzz_num = MAXFUZZNUMBER;
    if (argc > 1) { fuzz_num = atoi(argv[2]); }
    if (argc == 4) { deviceId = atoi(argv[3]); }
    printf("fuzz_num %d \n", fuzz_num);
    printf("aclnnAddLayerNormQuantV2_Normal :%d \n", aclnnAddLayerNormQuantV2_Normal(fuzz_num));
    return 0;
}