/*========================== begin_copyright_notice ============================

Copyright (C) 2017-2024 Intel Corporation

SPDX-License-Identifier: MIT

============================= end_copyright_notice ===========================*/

#pragma once

#include "ocl_igc_interface/igc_ocl_translation_ctx.h"
#include "ocl_igc_interface/impl/igc_ocl_device_ctx_impl.h"

#include <iomanip>
#include <memory>
#include <setjmp.h>
#include <signal.h>

#include "cif/builtins/memory/buffer/impl/buffer_impl.h"
#include "cif/helpers/error.h"
#include "cif/export/pimpl_base.h"

#include "ocl_igc_interface/impl/ocl_translation_output_impl.h"

#include "AdaptorOCL/OCL/TB/igc_tb.h"
#include "common/debug/Debug.hpp"

#include "cif/macros/enable.h"
#include <spirv-tools/libspirv.h>

#include "Compiler/CISACodeGen/OpenCLKernelCodeGen.hpp"
#include "OCLAPI/oclapi.h"

namespace TC{

// Taken from dllInterfaceCompute
extern bool ProcessElfInput(
  STB_TranslateInputArgs &InputArgs,
  STB_TranslateOutputArgs &OutputArgs,
  IGC::OpenCLProgramContext &Context,
  PLATFORM &platform,
  const TB_DATA_FORMAT& outType,
  float profilingTimerResolution);

extern bool ParseInput(
  llvm::Module*& pKernelModule,
  const STB_TranslateInputArgs* pInputArgs,
  STB_TranslateOutputArgs* pOutputArgs,
  IGC::OpenCLProgramContext &oclContext,
  TB_DATA_FORMAT inputDataFormatTemp);

OCL_API_CALL bool TranslateBuild(
  const STB_TranslateInputArgs* pInputArgs,
  STB_TranslateOutputArgs* pOutputArgs,
  TB_DATA_FORMAT inputDataFormatTemp,
  const IGC::CPlatform &platform,
  float profilingTimerResolution);

OCL_API_CALL bool TranslateBuildSPMD(
  const STB_TranslateInputArgs* pInputArgs,
  STB_TranslateOutputArgs* pOutputArgs,
  TB_DATA_FORMAT inputDataFormatTemp,
  const IGC::CPlatform &platform,
  float profilingTimerResolution,
  const ShaderHash& inputShHash);

OCL_API_CALL bool TranslateBuildVC(
  const STB_TranslateInputArgs* pInputArgs,
  STB_TranslateOutputArgs* pOutputArgs,
  TB_DATA_FORMAT inputDataFormatTemp,
  const IGC::CPlatform &platform,
  float profilingTimerResolution,
  const ShaderHash& inputShHash);

OCL_API_CALL void RebuildGlobalAnnotations(
  IGC::OpenCLProgramContext& oclContext,
  llvm::Module* pKernelModule
);

OCL_API_CALL bool ReadSpecConstantsFromSPIRV(
    std::istream &IS,
    std::vector<std::pair<uint32_t, uint32_t>> &OutSCInfo);

OCL_API_CALL void DumpShaderFile(
    const std::string& dstDir,
    const char* pBuffer,
    const UINT bufferSize,
    const QWORD hash,
    const std::string& ext,
    std::string* fileName);

OCL_API_CALL spv_result_t DisassembleSPIRV(
    const char* pBuffer,
    UINT bufferSize,
    spv_text* outSpirvAsm);

OCL_API_CALL void UnlockMutex();

}

OCL_API_CALL bool enableSrcLine(void*);

namespace IGC {

    OCL_API_CALL inline TC::TB_DATA_FORMAT toLegacyFormat(CodeType::CodeType_t format){
    switch(format){
        case CodeType::elf : return TC::TB_DATA_FORMAT_ELF; break;
        case CodeType::llvmBc : return TC::TB_DATA_FORMAT_LLVM_BINARY; break;
        case CodeType::llvmLl : return TC::TB_DATA_FORMAT_LLVM_TEXT; break;
        case CodeType::oclC : return TC::TB_DATA_FORMAT_OCL_TEXT; break;
        case CodeType::oclGenBin : return TC::TB_DATA_FORMAT_DEVICE_BINARY; break;
        case CodeType::spirV : return TC::TB_DATA_FORMAT_SPIR_V; break;
        default:
            return TC::TB_DATA_FORMAT_UNKNOWN;
    }
}

struct OCL_API_CALL IGC_State {
protected:
    bool isAlive;
public:
    IGC_State() {
        isAlive = true;
    }
    ~IGC_State() {
        isAlive = false;
    }
    IGC_State(const IGC_State&) = delete;
    IGC_State& operator=(const IGC_State&) = delete;
    static const bool isDestructed() {
        static const IGC_State obj;
        return !obj.isAlive;
    }
};

CIF_DECLARE_INTERFACE_PIMPL(IgcOclTranslationCtx) : CIF::PimplBase
{
    OCL_API_CALL CIF_PIMPL_DECLARE_CONSTRUCTOR(CIF::Version_t version, CIF_PIMPL(IgcOclDeviceCtx) *globalState,
                                  CodeType::CodeType_t inType, CodeType::CodeType_t outType)
        : globalState(CIF::Sanity::ToReferenceOrAbort(globalState)), inType(inType), outType(outType)
    {
    }

    OCL_API_CALL static bool SupportsTranslation(CodeType::CodeType_t inType, CodeType::CodeType_t outType){
        static std::pair<CodeType::CodeType_t, CodeType::CodeType_t> supportedTranslations[] =
            {
                  // from                 // to
                { CodeType::elf,      CodeType::llvmBc },
                { CodeType::elf,      CodeType::oclGenBin },
                { CodeType::llvmLl,   CodeType::oclGenBin },
                { CodeType::llvmBc,   CodeType::oclGenBin },
                { CodeType::spirV,    CodeType::oclGenBin },
            };
        for(const auto & st : supportedTranslations){
            if((inType == st.first) && (outType == st.second)){
                return true;
            }
        }
        return false;
    }

    OCL_API_CALL bool GetSpecConstantsInfo(CIF::Builtins::BufferSimple *src,
                              CIF::Builtins::BufferSimple *outSpecConstantsIds,
                              CIF::Builtins::BufferSimple *outSpecConstantsSizes)
    {
        if (IGC_State::isDestructed()) {
            return false;
        }
        bool success = false;
        const char* pInput = src->GetMemory<char>();
        uint32_t inputSize = static_cast<uint32_t>(src->GetSizeRaw());

        if(this->inType == CodeType::spirV){
            // load registry keys to make sure that flags can be read correctly
            LoadRegistryKeys();
            if (IGC_IS_FLAG_ENABLED(ShaderDumpEnable))
            {
                // dump SPIRV input if flag is enabled
                const char* pOutputFolder = Debug::GetShaderOutputFolder();
                QWORD hash = Debug::ShaderHashOCL(reinterpret_cast<const UINT*>(pInput), inputSize / 4).getAsmHash();

                TC::DumpShaderFile(pOutputFolder, pInput, inputSize, hash, ".spv", nullptr);
#if defined(IGC_SPIRV_TOOLS_ENABLED)
                spv_text spirvAsm = nullptr;
                if (TC::DisassembleSPIRV(pInput, inputSize, &spirvAsm) == SPV_SUCCESS)
                {
                    TC::DumpShaderFile(pOutputFolder, spirvAsm->str, spirvAsm->length, hash, ".spvasm", nullptr);
                }
                spvTextDestroy(spirvAsm);
#endif // defined(IGC_SPIRV_TOOLS_ENABLED)
            }
            llvm::StringRef strInput = llvm::StringRef(pInput, inputSize);
            std::istringstream IS(strInput.str());

            // vector of pairs [spec_id, spec_size]
            std::vector<std::pair<uint32_t, uint32_t>> SCInfo;
            success = TC::ReadSpecConstantsFromSPIRV(IS, SCInfo);

            outSpecConstantsIds->Resize(sizeof(uint32_t) * SCInfo.size());
            outSpecConstantsSizes->Resize(sizeof(uint32_t) * SCInfo.size());

            uint32_t* specConstantsIds = outSpecConstantsIds->GetMemoryWriteable<uint32_t>();
            uint32_t* specConstantsSizes = outSpecConstantsSizes->GetMemoryWriteable<uint32_t>();

            for(uint32_t i = 0; i < SCInfo.size(); ++i){
                specConstantsIds[i] = SCInfo.at(i).first;
                specConstantsSizes[i] = SCInfo.at(i).second;
            }
        }
        else{
            success = false;
        }

        return success;
    }

    OCL_API_CALL OclTranslationOutputBase *Translate(CIF::Version_t outVersion,
                                        CIF::Builtins::BufferSimple *src,
                                        CIF::Builtins::BufferSimple *specConstantsIds,
                                        CIF::Builtins::BufferSimple *specConstantsValues,
                                        CIF::Builtins::BufferSimple *options,
                                        CIF::Builtins::BufferSimple *internalOptions,
                                        CIF::Builtins::BufferSimple *tracingOptions,
                                        uint32_t tracingOptionsCount,
                                        void *gtPinInput) const{
        // Create interface for return data
        auto outputInterface = CIF::RAII::UPtr(CIF::InterfaceCreator<OclTranslationOutput>::CreateInterfaceVer(outVersion, this->outType));
        if(outputInterface == nullptr){
            return nullptr; // OOM
        }
        if (IGC_State::isDestructed()) {
            outputInterface->GetImpl()->SetError(TranslationErrorType::UnhandledInput, "IGC is destructed");
            return outputInterface.release();
        }

        TC::STB_TranslateInputArgs inputArgs;
        if(src != nullptr){
            if (gtPinInput)
            {
                bool srcLine = enableSrcLine(gtPinInput);
                if (srcLine)
                {
                    const char* arg = " -gline-tables-only";
                    src->PushBackRawBytes(arg, strlen(arg));
                }
            }
            inputArgs.pInput = src->GetMemoryWriteable<char>();
            inputArgs.InputSize = static_cast<uint32_t>(src->GetSizeRaw());
        }
        if(options != nullptr){
            inputArgs.pOptions = options->GetMemory<char>();
            inputArgs.OptionsSize = static_cast<uint32_t>(options->GetSizeRaw());
        }
        if(internalOptions != nullptr){
            inputArgs.pInternalOptions =  internalOptions->GetMemory<char>();
            inputArgs.InternalOptionsSize = static_cast<uint32_t>(internalOptions->GetSizeRaw());
        }
        if(tracingOptions != nullptr){
            inputArgs.pTracingOptions = tracingOptions->GetMemoryRawWriteable();
        }
        inputArgs.TracingOptionsCount = tracingOptionsCount;
        if(specConstantsIds != nullptr && specConstantsValues != nullptr){
            inputArgs.pSpecConstantsIds = specConstantsIds->GetMemory<uint32_t>();
            inputArgs.SpecConstantsSize = static_cast<uint32_t>(specConstantsIds->GetSizeRaw() / sizeof(uint32_t));
            inputArgs.pSpecConstantsValues = specConstantsValues->GetMemory<uint64_t>();
        }
        inputArgs.GTPinInput = gtPinInput;

        CIF::Sanity::NotNullOrAbort(this->globalState.GetPlatformImpl());
        auto platform = this->globalState.GetPlatformImpl()->p;

        USC::SShaderStageBTLayout zeroLayout = USC::g_cZeroShaderStageBTLayout;
        IGC::COCLBTILayout oclLayout(&zeroLayout);

        TC::STB_TranslateOutputArgs output;
        CIF::SafeZeroOut(output);

        std::string RegKeysFlagsFromOptions;
        if (inputArgs.pOptions != nullptr)
        {
            std::string_view optionsWithFlags = inputArgs.pOptions;

            // check if there are more instances of '-igc_opts'
            while (!optionsWithFlags.empty())
            {
                std::size_t found = optionsWithFlags.find("-igc_opts");
                if (found == std::string::npos)
                    break;

                std::size_t foundFirstSingleQuote = optionsWithFlags.find('\'', found);
                std::size_t foundSecondSingleQuote = optionsWithFlags.find('\'', foundFirstSingleQuote + 1);
                if (foundFirstSingleQuote != std::string::npos && foundSecondSingleQuote)
                {
                    RegKeysFlagsFromOptions += optionsWithFlags.substr(foundFirstSingleQuote + 1, (foundSecondSingleQuote - foundFirstSingleQuote - 1));
                    RegKeysFlagsFromOptions = RegKeysFlagsFromOptions + ',';
                    optionsWithFlags = optionsWithFlags.substr(foundSecondSingleQuote + 1, optionsWithFlags.length());
                }
            }
        }
        bool RegFlagNameError = 0;
        LoadRegistryKeys(RegKeysFlagsFromOptions, &RegFlagNameError);
        if(RegFlagNameError) outputInterface->GetImpl()->SetError(TranslationErrorType::Unused, "Invalid registry flag name in -igc_opts, at least one flag has been ignored");

        IGC::CPlatform igcPlatform = this->globalState.GetIgcCPlatform();

        // extra ocl options set from regkey
        const char *extraOptions = IGC_GET_REGKEYSTRING(ExtraOCLOptions);
        std::string combinedOptions;
        if (extraOptions[0] != '\0')
        {
            if (inputArgs.pOptions != nullptr)
            {
                combinedOptions = std::string(inputArgs.pOptions) + ' ';
            }
            combinedOptions += extraOptions;
            inputArgs.pOptions = combinedOptions.c_str();
            inputArgs.OptionsSize = combinedOptions.size();
        }

        // extra ocl internal options set from regkey
        const char *extraInternlOptions = IGC_GET_REGKEYSTRING(ExtraOCLInternalOptions);
        std::string combinedInternalOptions;
        if (extraInternlOptions[0] != '\0')
        {
            if (inputArgs.pInternalOptions != nullptr)
            {
                combinedInternalOptions = std::string(inputArgs.pInternalOptions) + ' ';
            }
            combinedInternalOptions += extraInternlOptions;
            inputArgs.pInternalOptions = combinedInternalOptions.c_str();
            inputArgs.InternalOptionsSize = combinedInternalOptions.size();
        }

        bool success = false;
        try
        {
            if (this->inType == CodeType::elf)
            {
                // Handle TB_DATA_FORMAT_ELF input as a result of a call to
                // clLinkLibrary(). There are two possible scenarios, link input
                // to form a new library (BC module) or link input to form an
                // executable.

                // First, link input modules together
                CDriverInfo dummyDriverInfo;
                IGC::OpenCLProgramContext oclContextTemp(oclLayout, igcPlatform, &inputArgs, dummyDriverInfo, nullptr, false);
                IGC::Debug::RegisterComputeErrHandlers(*oclContextTemp.getLLVMContext());
                success = TC::ProcessElfInput(
                    inputArgs, output, oclContextTemp, platform,
                    toLegacyFormat(this->outType),
                    this->globalState.MiscOptions.ProfilingTimerResolution);
            }
            else
            {
                if ((this->inType == CodeType::llvmLl) ||
                    (this->inType == CodeType::spirV) ||
                    (this->inType == CodeType::llvmBc))
                {
                    TC::TB_DATA_FORMAT inFormatLegacy = toLegacyFormat(this->inType);
                    success = TC::TranslateBuild(
                        &inputArgs,
                        &output,
                        inFormatLegacy,
                        igcPlatform,
                        this->globalState.MiscOptions.ProfilingTimerResolution);
                }
                else
                {
                    outputInterface->GetImpl()->SetError(TranslationErrorType::UnhandledInput, "Unhandled inType");
                    success = false;
                }
            }
        }
        catch (std::exception& e)
        {
            success = false;
            if (output.ErrorStringSize == 0 && output.pErrorString == nullptr)
            {
                std::string msg = "IGC: ";
                msg += e.what();
                outputInterface->GetImpl()->SetError(TranslationErrorType::FailedCompilation, msg.c_str());
            }
        }
        catch (...)
        {
            success = false;
            if (output.ErrorStringSize == 0 && output.pErrorString == nullptr)
            {
                outputInterface->GetImpl()->SetError(TranslationErrorType::FailedCompilation, "IGC: Internal Compiler Error");
            }
        }

        auto outputData = std::unique_ptr<char[]>(output.pOutput);
        auto errorString = std::unique_ptr<char[]>(output.pErrorString);
        auto debugData = std::unique_ptr<char[]>(output.pDebugData);

        bool dataCopiedSuccessfuly = true;
        if(success){
            dataCopiedSuccessfuly &= outputInterface->GetImpl()->AddWarning(output.pErrorString, output.ErrorStringSize);
            dataCopiedSuccessfuly &= outputInterface->GetImpl()->CloneDebugData(output.pDebugData, output.DebugDataSize);
            dataCopiedSuccessfuly &= outputInterface->GetImpl()->SetSuccessfulAndCloneOutput(output.pOutput, output.OutputSize);
        }else{
            dataCopiedSuccessfuly &= outputInterface->GetImpl()->SetError(TranslationErrorType::FailedCompilation, output.pErrorString);
        }

        if(dataCopiedSuccessfuly == false){
            return nullptr; // OOM
        }

        return outputInterface.release();
    }

    OCL_API_CALL OclTranslationOutputBase* GetErrorOutput(CIF::Version_t outVersion, unsigned int code) const
    {
        auto outputInterface = CIF::RAII::UPtr(CIF::InterfaceCreator<OclTranslationOutput>::CreateInterfaceVer(outVersion, this->outType));
        if (outputInterface == nullptr) {
            return nullptr; // OOM
        }
        const char* outputMessage;
        switch (code) {
#if defined(WIN32)
        case EXCEPTION_ACCESS_VIOLATION:
            outputMessage = "IGC: Internal Compiler Error: Access violation";
            break;
        case EXCEPTION_DATATYPE_MISALIGNMENT:
            outputMessage = "IGC: Internal Compiler Error: Datatype misalignement";
            break;
        case EXCEPTION_FLT_DIVIDE_BY_ZERO:
        case EXCEPTION_INT_DIVIDE_BY_ZERO:
            outputMessage = "IGC: Internal Compiler Error: Divide by zero";
            break;
        case EXCEPTION_STACK_OVERFLOW:
            outputMessage = "IGC: Internal Compiler Error: Stack overflow";
            break;
#else
        case SIGABRT:
            outputMessage = "IGC: Internal Compiler Error: Abnormal termination";
            break;
        case SIGFPE:
            outputMessage = "IGC: Internal Compiler Error: Floating point exception";
            break;
        case SIGILL:
            outputMessage = "IGC: Internal Compiler Error: Invalid instruction";
            break;
        case SIGINT:
            outputMessage = "IGC: Internal Compiler Error: Interrupt request sent to the program";
            break;
        case SIGSEGV:
            outputMessage = "IGC: Internal Compiler Error: Segmentation violation";
            break;
        case SIGTERM:
            outputMessage = "IGC: Internal Compiler Error: Termination request sent to the program";
            break;
#endif
        default:
            outputMessage = "IGC: Internal Compiler Error: Signal caught";
        }
        if (outputInterface->GetImpl()->SetError(TranslationErrorType::Internal, outputMessage)) {
            return outputInterface.release();
        }
        return nullptr;
    }

protected:
    CIF_PIMPL(IgcOclDeviceCtx) &globalState;
    CodeType::CodeType_t inType;
    CodeType::CodeType_t outType;
};

CIF_DEFINE_INTERFACE_TO_PIMPL_FORWARDING_CTOR_DTOR(IgcOclTranslationCtx);

}

#include "cif/macros/disable.h"

#if defined(NDEBUG)
#if defined(WIN32)
OCL_API_CALL int ex_filter(unsigned int code, struct _EXCEPTION_POINTERS* ep);

#define EX_GUARD_BEGIN                                                        \
__try                                                                         \
{                                                                             \

#define EX_GUARD_END                                                          \
}                                                                             \
__except (ex_filter(GetExceptionCode(), GetExceptionInformation()))           \
{                                                                             \
    TC::UnlockMutex();                                                        \
    res = CIF_GET_PIMPL()->GetErrorOutput(outVersion, GetExceptionCode());    \
}                                                                             \

#else
OCL_API_CALL void signalHandler(int sig, siginfo_t* info, void* ucontext);

#include "igc_signal_guard.h"

#define EX_GUARD_BEGIN                                                         \
  do {                                                                         \
    SET_SIG_HANDLER(SIGABRT)                                                   \
    SET_SIG_HANDLER(SIGFPE)                                                    \
    SET_SIG_HANDLER(SIGILL)                                                    \
    SET_SIG_HANDLER(SIGINT)                                                    \
    SET_SIG_HANDLER(SIGSEGV)                                                   \
    SET_SIG_HANDLER(SIGTERM)                                                   \
    int sig = setjmp(sig_jmp_buf);                                             \
    if (sig == 0) {

#define EX_GUARD_END                                                           \
    } else {                                                                   \
      TC::UnlockMutex();                                                       \
      res = CIF_GET_PIMPL()->GetErrorOutput(outVersion, sig);                  \
    }                                                                          \
    REMOVE_SIG_HANDLER(SIGABRT)                                                \
    REMOVE_SIG_HANDLER(SIGFPE)                                                 \
    REMOVE_SIG_HANDLER(SIGILL)                                                 \
    REMOVE_SIG_HANDLER(SIGINT)                                                 \
    REMOVE_SIG_HANDLER(SIGSEGV)                                                \
    REMOVE_SIG_HANDLER(SIGTERM)                                                \
  } while (0);

#endif
#else
#define EX_GUARD_BEGIN
#define EX_GUARD_END
#endif
