#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_advance_step.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 aclnnAdvanceStep_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("aclnnAdvanceStep", "aclnnAdvanceStep op fuzz start \n");
    const long long kTotalElemsMax = 1048576LL;

    aclrtSetOpExecuteTimeOut(1091);
    DT_Set_TimeOut_Second(3000);
    DT_FUZZ_START(0, fuzz_num, "aclnnAdvanceStep", 0)
    {
        int i = 0;
        int aclnnRet = -1;
        void* itAddr = nullptr; aclTensor* inputTokens = nullptr;
        void* stAddr = nullptr; aclTensor* sampledTokenIds = nullptr;
        void* ipAddr = nullptr; aclTensor* inputPositions = nullptr;
        void* slAddr = nullptr; aclTensor* seqLens = nullptr;
        void* smAddr = nullptr; aclTensor* slotMapping = nullptr;
        void* btAddr = nullptr; aclTensor* blockTables = nullptr;
        uint64_t workspaceSize = 0;
        aclOpExecutor* executor = nullptr;
        void* workspaceAddr = nullptr;

        LOG_PRINT("aclnnAdvanceStep_Normal %d\n", fuzzi);
        printf("[FUZZ_LOG] plog_file: aclnnAdvanceStep_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 numSeqs = *(int *)DTSetGetNumberRange(&g_Element[i++], 2, 2, 64, "i");       // numSeqs ∈ [2,64]
            int numQueries = *(int *)DTSetGetNumberRange(&g_Element[i++], 2, 1, 63, "i");    // numQueries < numSeqs
            int blockSize = *(int *)DTSetGetNumberRange(&g_Element[i++], 2, 1, 16, "i");
            int seqLenMax = *(int *)DTSetGetNumberRange(&g_Element[i++], 2, 1, 64, "i");
            // size bucket 调整规模
            const int seqTargets[4] = {8, 64, 256, 1024};
            if (numSeqs > seqTargets[sizeBucket]) { numSeqs = seqTargets[sizeBucket]; }
            if (numQueries >= numSeqs) { numQueries = numSeqs - 1; }
            if (numQueries < 1) { numQueries = 1; }
            if (blockSize < 1) { blockSize = 1; }
            if (seqLenMax < 1) { seqLenMax = 1; }
            int64_t v_numSeqs = (int64_t)numSeqs;
            int64_t v_numQueries = (int64_t)numQueries;
            int64_t v_blockSize = (int64_t)blockSize;

            // shapes
            std::vector<int64_t> seqShape = {v_numSeqs};
            std::vector<int64_t> sampledShape = {v_numQueries, 1};
            int64_t maxSeqLen = (int64_t)seqLenMax;
            int64_t maxBlocks = maxSeqLen / v_blockSize + 1;  // 第二维 > max(seqLens)/blockSize
            std::vector<int64_t> btShape = {v_numSeqs, maxBlocks};

            long long totalElems = GetShapeSize(seqShape) * 4 + GetShapeSize(sampledShape) + GetShapeSize(btShape);
            if (totalElems > kTotalElemsMax) { printf("shape_size too big\n"); goto ERROR; }

            // create tensors（数据敏感：全部正整数）
            {
                std::vector<int64_t> itH((size_t)numSeqs);
                std::vector<int64_t> ipH((size_t)numSeqs);
                std::vector<int64_t> slH((size_t)numSeqs);
                std::vector<int64_t> smH((size_t)numSeqs);
                for (int k = 0; k < numSeqs; k++) {
                    itH[k] = (int64_t)k + 1;
                    ipH[k] = (int64_t)k + 1;
                    slH[k] = (int64_t)((k % seqLenMax) + 1);  // 正整数, ≤ seqLenMax
                    smH[k] = (int64_t)k + 1;
                }
                std::vector<int64_t> stH((size_t)numQueries);
                for (int k = 0; k < numQueries; k++) { stH[k] = (int64_t)k + 1; }
                std::vector<int64_t> btH((size_t)(numSeqs * maxBlocks));
                for (int64_t k = 0; k < (int64_t)(numSeqs * maxBlocks); k++) { btH[k] = k + 1; }
                ret = CreateAclTensor_New(itH, seqShape, &itAddr, ACL_INT64, aclFormat::ACL_FORMAT_ND, &inputTokens);
                CHECK_RET(ret == ACL_SUCCESS, goto ERROR);
                ret = CreateAclTensor_New(stH, sampledShape, &stAddr, ACL_INT64, aclFormat::ACL_FORMAT_ND, &sampledTokenIds);
                CHECK_RET(ret == ACL_SUCCESS, goto ERROR);
                ret = CreateAclTensor_New(ipH, seqShape, &ipAddr, ACL_INT64, aclFormat::ACL_FORMAT_ND, &inputPositions);
                CHECK_RET(ret == ACL_SUCCESS, goto ERROR);
                ret = CreateAclTensor_New(slH, seqShape, &slAddr, ACL_INT64, aclFormat::ACL_FORMAT_ND, &seqLens);
                CHECK_RET(ret == ACL_SUCCESS, goto ERROR);
                ret = CreateAclTensor_New(smH, seqShape, &smAddr, ACL_INT64, aclFormat::ACL_FORMAT_ND, &slotMapping);
                CHECK_RET(ret == ACL_SUCCESS, goto ERROR);
                ret = CreateAclTensor_New(btH, btShape, &btAddr, ACL_INT64, aclFormat::ACL_FORMAT_ND, &blockTables);
                CHECK_RET(ret == ACL_SUCCESS, goto ERROR);
            }
            ret = aclnnAdvanceStepGetWorkspaceSize(inputTokens, sampledTokenIds, inputPositions, seqLens,
                slotMapping, blockTables, v_numSeqs, v_numQueries, v_blockSize, &workspaceSize, &executor);
            printf("[FUZZ_PARAM] inputTokens: dtype=INT64, shape=%s; sampledTokenIds: shape=%s; blockTables: shape=%s; numSeqs=%lld, numQueries=%lld, blockSize=%lld, maxSeqLen=%lld\n",
                   ShapeToStr(seqShape).c_str(), ShapeToStr(sampledShape).c_str(), ShapeToStr(btShape).c_str(),
                   (long long)v_numSeqs, (long long)v_numQueries, (long long)v_blockSize, (long long)maxSeqLen);
            CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnAdvanceStepGetWorkspaceSize 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 = aclnnAdvanceStep(workspaceAddr, workspaceSize, executor, stream);
            aclnnRet = (int)ret;
            if (ret == ACL_SUCCESS) { printf("aclnnAdvanceStep execute success :%d\n", (int)ret); }
            else { printf("aclnnAdvanceStep 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("aclnnAdvanceStep 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 (inputTokens) { aclDestroyTensor(inputTokens); }
        if (sampledTokenIds) { aclDestroyTensor(sampledTokenIds); }
        if (inputPositions) { aclDestroyTensor(inputPositions); }
        if (seqLens) { aclDestroyTensor(seqLens); }
        if (slotMapping) { aclDestroyTensor(slotMapping); }
        if (blockTables) { aclDestroyTensor(blockTables); }
        if (itAddr) { aclrtFree(itAddr); }
        if (stAddr) { aclrtFree(stAddr); }
        if (ipAddr) { aclrtFree(ipAddr); }
        if (slAddr) { aclrtFree(slAddr); }
        if (smAddr) { aclrtFree(smAddr); }
        if (btAddr) { aclrtFree(btAddr); }
        if (workspaceSize > 0 && workspaceAddr) { aclrtFree(workspaceAddr); }
    }
    DT_FUZZ_END()
    WriteDTLog("aclnnAdvanceStep", "aclnnAdvanceStep 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("aclnnAdvanceStep", "aclnnAdvanceStep op fuzz end \n");
    aclrtDestroyStream(stream);
    aclrtDestroyContext(context);
    aclrtResetDevice(deviceId);
    aclFinalize();
    return 0;
}

extern "C" int aclnnAdvanceStep_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("aclnnAdvanceStep_Normal :%d \n", aclnnAdvanceStep_Normal(fuzz_num));
    return 0;
}