#include "iwa/util/glsl_compiler.hpp"

#include <filesystem>
#include <utility>
#include <glslang/Include/InfoSink.h>
#include <glslang/Public/ShaderLang.h>
#include <glslang/MachineIndependent/iomapper.h>
#include <glslang/MachineIndependent/localintermediate.h>
#include <glslang/Public/ResourceLimits.h>
#include <glslang/SPIRV/GlslangToSpv.h>
#include <yaml-cpp/yaml.h>
#include "iwa/device.hpp"
#include "iwa/instance.hpp"
#include "iwa/log.hpp"
#include "iwa/util/dir_stack_file_includer.hpp"
#include "iwa/util/reflect_glsl.hpp"

namespace fs = std::filesystem;

namespace iwa
{
namespace
{
class CustomFileIncluder : public impl::DirStackFileIncluder
{
private:
    fs::path workingDir;
    mijin::FileSystemAdapter& mFsAdapter;
public:
    explicit CustomFileIncluder(mijin::FileSystemAdapter& fsAdapter) noexcept: mFsAdapter(fsAdapter)
    {}

public:
    void setWorkingDir(const fs::path& workingDir_)
    { workingDir = workingDir_; }

protected: // overrides
    IncludeResult* readLocalPath(const char* headerName, const char* includerName, int depth) override
    {
        // Discard popped include directories, and
        // initialize when at parse-time first level.
        directoryStack.resize(depth + externalLocalDirectoryCount);
        if (depth == 1)
        {
            directoryStack.back() = getDirectory(includerName);
        }

        // Find a directory that works, using a reverse search of the include stack.
        for (auto it = directoryStack.rbegin(); it != directoryStack.rend(); ++it)
        {
            std::string path = *it + '/' + headerName;
            std::replace(path.begin(), path.end(), '\\', '/');

            std::unique_ptr<mijin::Stream> stream;
            mijin::StreamError error = mijin::StreamError::UNKNOWN_ERROR;
            if (workingDir != fs::path())
            {
                // try relative include first
                error = mFsAdapter.open(workingDir / path, mijin::FileOpenMode::READ, stream);
            }
            if (error != mijin::StreamError::SUCCESS)
            {
                error = mFsAdapter.open(path, mijin::FileOpenMode::READ, stream);
            }
            if (error == mijin::StreamError::SUCCESS)
            {
                directoryStack.push_back(getDirectory(path));
                includedFiles.insert(path);
                return newCustomIncludeResult(path, *stream);
            }
        }

        return nullptr;
    }

    // Do actual reading of the file, filling in a new include result.
    IncludeResult* newCustomIncludeResult(const std::string& path, mijin::Stream& stream) const
    {
        (void) stream.seek(0, mijin::SeekMode::RELATIVE_TO_END);
        const std::size_t length = stream.tell();
        (void) stream.seek(0);

        char* content = new tUserDataElement[length]; // NOLINT(cppcoreguidelines-owning-memory)
        const mijin::StreamError error = stream.readRaw(content, length);
        if (error != mijin::StreamError::SUCCESS)
        {
            logAndDie("Error reading include file.");
        }
        return new IncludeResult(path, content, length, content); // NOLINT(cppcoreguidelines-owning-memory)
    }
};

class SemanticIoResolver : public glslang::TDefaultIoResolverBase
{
private:
    const std::vector<GLSLSemanticMapping>& mMappings;
public:
    SemanticIoResolver(const glslang::TIntermediate& intermediate, const std::vector<GLSLSemanticMapping>& mappings)
        : glslang::TDefaultIoResolverBase(intermediate), mMappings(mappings) {}

    bool validateBinding(EShLanguage /* stage */, glslang::TVarEntryInfo& /* ent */) override { return true; }

    glslang::TResourceType getResourceType(const glslang::TType& type) override {
        if (isImageType(type)) {
            return glslang::EResImage;
        }
        if (isTextureType(type)) {
            return glslang::EResTexture;
        }
        if (isSsboType(type)) {
            return glslang::EResSsbo;
        }
        if (isSamplerType(type)) {
            return glslang::EResSampler;
        }
        if (isUboType(type)) {
            return glslang::EResUbo;
        }
        return glslang::EResCount;
    }

    int resolveBinding(EShLanguage stage, glslang::TVarEntryInfo& ent) override
    {
        const glslang::TType& type = ent.symbol->getType();
        if (type.getQualifier().hasSemantic())
        {
            const unsigned semantic = type.getQualifier().layoutSemantic;
            const unsigned semanticIdx = type.getQualifier().hasSemanticIndex() ? type.getQualifier().layoutSemanticIndex : 0;
            auto it = std::ranges::find_if(mMappings, [&](const GLSLSemanticMapping& mapping)
            {
                return mapping.semantic == semantic && mapping.semanticIdx == semanticIdx;
            });
            if (it != mMappings.end()) {
                return ent.newBinding = it->newBinding;
            }
        }

        // default implementation
        const int set = getLayoutSet(type);
        // On OpenGL arrays of opaque types take a seperate binding for each element
        const int numBindings = referenceIntermediate.getSpv().openGl != 0 && type.isSizedArray() ? type.getCumulativeArraySize() : 1;
        const glslang::TResourceType resource = getResourceType(type);
        if (resource < glslang::EResCount) {
            if (type.getQualifier().hasBinding()) {
                return ent.newBinding = reserveSlot(
                    set, getBaseBinding(stage, resource, set) + type.getQualifier().layoutBinding, numBindings);
            }
            if (ent.live && doAutoBindingMapping()) {
                // find free slot, the caller did make sure it passes all vars with binding
                // first and now all are passed that do not have a binding and needs one
                return ent.newBinding = getFreeSlot(set, getBaseBinding(stage, resource, set), numBindings);
            }
        }
        return ent.newBinding = -1;
    }
    int resolveSet(EShLanguage stage, glslang::TVarEntryInfo& ent) override
    {
        const glslang::TType& type = ent.symbol->getType();
        if (type.getQualifier().hasSemantic())
        {
            const unsigned semantic = type.getQualifier().layoutSemantic;
            const unsigned semanticIdx = type.getQualifier().hasSemanticIndex() ? type.getQualifier().layoutSemanticIndex : 0;
            auto it = std::ranges::find_if(mMappings, [&](const GLSLSemanticMapping& mapping)
            {
                return mapping.semantic == semantic && mapping.semanticIdx == semanticIdx;
            });
            if (it == mMappings.end()) {
                return glslang::TDefaultIoResolverBase::resolveSet(stage, ent);
            }
            return ent.newSet = it->newSet;
        }
        return glslang::TDefaultIoResolverBase::resolveSet(stage, ent);
    }
    void addStage(EShLanguage stage, glslang::TIntermediate& stageIntermediate) override
    {
        nextInputLocation = nextOutputLocation = 0;
        glslang::TDefaultIoResolverBase::addStage(stage, stageIntermediate);
    }
};

void initGlslang() noexcept
{
    static bool inited = false;
    if (inited)
    {
        return;
    }
    inited = true;
    if (!glslang::InitializeProcess())
    {
        logAndDie("Error initializing Glslang.");
    }
}
} // namespace

EShLanguage typeToGlslang(vk::ShaderStageFlagBits type) noexcept
{
    switch (type)
    {
        case vk::ShaderStageFlagBits::eCompute:
            return EShLangCompute;
        case vk::ShaderStageFlagBits::eVertex:
            return EShLangVertex;
        case vk::ShaderStageFlagBits::eFragment:
            return EShLangFragment;
        case vk::ShaderStageFlagBits::eRaygenKHR:
            return EShLangRayGen;
        case vk::ShaderStageFlagBits::eClosestHitKHR:
            return EShLangClosestHit;
        case vk::ShaderStageFlagBits::eAnyHitKHR:
            return EShLangAnyHit;
        case vk::ShaderStageFlagBits::eMissKHR:
            return EShLangMiss;
        case vk::ShaderStageFlagBits::eIntersectionKHR:
            return EShLangIntersect;
        case vk::ShaderStageFlagBits::eCallableKHR:
            return EShLangCallable;
        case vk::ShaderStageFlagBits::eTaskEXT:
            return EShLangTask;
        case vk::ShaderStageFlagBits::eMeshEXT:
            return EShLangMesh;
        case vk::ShaderStageFlagBits::eTessellationControl:
            return EShLangTessControl;
        case vk::ShaderStageFlagBits::eTessellationEvaluation:
            return EShLangTessEvaluation;
        case vk::ShaderStageFlagBits::eGeometry:
            return EShLangGeometry;
        case vk::ShaderStageFlagBits::eAllGraphics:
        case vk::ShaderStageFlagBits::eAll:
        case vk::ShaderStageFlagBits::eSubpassShadingHUAWEI:
        case vk::ShaderStageFlagBits::eClusterCullingHUAWEI:
            break; // let it fail
    }

    logAndDie("Invalid value passed to typeToGlslang!");
}

ShaderSource ShaderSource::fromStream(mijin::Stream& stream, std::string fileName)
{
    ShaderSource result = {
        .fileName = std::move(fileName)
    };
    if (const mijin::StreamError error = stream.readAsString(result.code); error != mijin::StreamError::SUCCESS)
    {
        // TODO: custom exception type, for stacktrace and stuff
        throw std::runtime_error("Error reading shader source.");
    }
    return result;
}

ShaderSource ShaderSource::fromFile(const mijin::PathReference& file)
{
    std::unique_ptr<mijin::Stream> stream;
    if (const mijin::StreamError error = file.open(mijin::FileOpenMode::READ, stream); error != mijin::StreamError::SUCCESS)
    {
        throw std::runtime_error("Error opening file for reading shader source.");
    }
    return fromStream(*stream, file.getPath().string());
}

ShaderSource ShaderSource::fromYaml(const YAML::Node& node, const mijin::PathReference& yamlFile)
{
    if (node.Tag() == "!load")
    {
        return fromFile(yamlFile.getAdapter()->getPath(node.as<std::string>()));
    }
    const std::string& source = node["source"].as<std::string>();
    std::string fileName;
    if (const YAML::Node fileNameNode = node["fileName"]; !fileNameNode.IsNull())
    {
        fileName = fileNameNode.as<std::string>();
    }
    return {
        .code = source,
        .fileName = std::move(fileName)
    };
}

GLSLShader::GLSLShader(ObjectPtr<Instance> owner, GLSLShaderCreationArgs args) noexcept
        : super_t(std::move(owner)), mType(args.type), mSources(std::move(args.sources)), mDefines(std::move(args.defines))
{
    MIJIN_ASSERT(!mSources.empty(), "Cannot compile without sources.");
    compile();
}

GLSLShader::~GLSLShader() noexcept = default;

std::unique_ptr<glslang::TShader> GLSLShader::releaseHandle()
{
    if (mHandle == nullptr) {
        compile();
    }
    return std::exchange(mHandle, nullptr);
}

ShaderMeta GLSLShader::getPartialMeta()
{
    if (mHandle == nullptr) {
        compile();
    }
    return reflectShader(*mHandle);
}

void GLSLShader::compile()
{
    initGlslang();

    const EShLanguage stage = typeToGlslang(mType);

    std::unique_ptr<glslang::TShader> shader = std::make_unique<glslang::TShader>(stage); // NOLINT(cppcoreguidelines-owning-memory)

    std::vector<const char*> sources;
    std::vector<int> lengths;
    std::vector<const char*> names;
    sources.reserve(mSources.size() + 1);
    lengths.reserve(mSources.size() + 1);
    names.reserve(mSources.size() + 1);

    std::string preamble = getOwner()->getInstanceExtension<GLSLCompilerSettings>().getCommonPreamble();
    for (const std::string& define : mDefines) {
        preamble.append(fmt::format("\n#define {}\n", define));
    }
    sources.push_back(preamble.c_str());
    lengths.push_back(static_cast<int>(preamble.size()));
    names.push_back("<preamble>");

    for (const ShaderSource& source : mSources)
    {
        sources.push_back(source.code.c_str());
        lengths.push_back(static_cast<int>(source.code.size()));
        names.push_back(source.fileName.c_str());
    }
    shader->setStringsWithLengthsAndNames(sources.data(), lengths.data(), names.data(), static_cast<int>(sources.size()));
    shader->setDebugInfo(true);
    shader->setEnvInput(glslang::EShSourceGlsl, stage, glslang::EShClientVulkan, glslang::EShTargetVulkan_1_3);
    shader->setEnvClient(glslang::EShClientVulkan, glslang::EShTargetVulkan_1_3);
    shader->setEnvTarget(glslang::EShTargetLanguage::EShTargetSpv, glslang::EShTargetSpv_1_6);
    shader->setAutoMapLocations(true);
    shader->setAutoMapBindings(true);

    const EShMessages PREPROCESS_MESSAGES = static_cast<EShMessages>(EShMsgDefault
#if !defined(KAZAN_RELEASE)
        | EShMsgDebugInfo
#endif
    );
    std::string completeCode;
    CustomFileIncluder includer(getOwner()->getPrimaryFSAdapter()); // TODO: this type seems to be doing stupid things, investigate
    const std::string sourceFileAbsStr = mSources[0].fileName; // just for you MSVC <3
    if (!sourceFileAbsStr.empty()) {
        includer.setWorkingDir(fs::path(sourceFileAbsStr).parent_path());
    }
    const bool couldPreprocess = shader->preprocess(
        /* builtInResources              = */ GetDefaultResources(),
        /* defaultVersion                = */ 450,
        /* defaultProfile                = */ ECoreProfile,
        /* forceDefaultVersionAndProfile = */ false,
        /* forwardCompatible             = */ false,
        /* message                       = */ PREPROCESS_MESSAGES,
        /* outputString                  = */ &completeCode,
        /* includer                      = */ includer
    );
    if (!couldPreprocess)
    {
        logMsg("GLSL preprocessing failed:\ninfo log:\n{}\ndebug log:\n{}",
               shader->getInfoLog(), shader->getInfoDebugLog()
        );
        logAndDie("Error preprocessing shader.");
    }

#if 0
    ShaderPreprocessResult preprocessResult = preprocessShader(completeCode);

    for (std::string& module : preprocessResult.importedModules) {
        importedModules.insert(std::move(module));
    }

    for (std::string& option : preprocessResult.supportedOptions) {
        supportedOptions.insert(std::move(option));
    }
#endif

    const char* newSource = completeCode.c_str();
#if defined(KAZAN_RELEASE)
    shader->setStrings(&newSource, 1); // replace source with the preprocessed one
#else
    const int newSourceLen = static_cast<int>(std::strlen(newSource));
    const char* newSourceName = sourceFileAbsStr.c_str();
    shader->setStringsWithLengthsAndNames(&newSource, &newSourceLen, &newSourceName, 1);
#endif
    const EShMessages PARSE_MESSAGES = static_cast<EShMessages>(EShMsgDefault
#if !defined(KAZAN_RELEASE)
        | EShMsgDebugInfo
#endif
    );
    const bool couldParse = shader->parse(
        /* builtinResources  = */ GetDefaultResources(),
        /* defaultVersion    = */ 450,
        /* forwardCompatible = */ false,
        /* messages          = */ PARSE_MESSAGES
    );
    if (!couldParse)
    {
        logMsg("GLSL parsing failed:\ninfo log:\n{}\ndebug log:\n{}",
               shader->getInfoLog(), shader->getInfoDebugLog()
        );
        logAndDie("Error parsing shader.");
    }

    mHandle = std::move(shader);
}

GLSLShaderProgram::GLSLShaderProgram(ObjectPtr<Device> owner, GLSLShaderProgramCreationArgs args)
    : super_t(std::move(owner)), mLinkFlags(args.linkFlags)
{
    MIJIN_ASSERT_FATAL(!args.shaders.empty(), "At least one shader per program is required!");

    mHandle = std::make_unique<glslang::TProgram>();
    for (const ObjectPtr<GLSLShader>& shader : args.shaders)
    {
        mShaderHandles.push_back(shader->releaseHandle());
        mHandle->addShader(mShaderHandles.back().get());
    }

    const EShMessages linkMessages = static_cast<EShMessages>(EShMsgSpvRules | EShMsgVulkanRules
                                   | (args.linkFlags.withDebugInfo ? EShMsgDebugInfo : EShMessages(0)));

    if (!mHandle->link(linkMessages))
    {
        logAndDie("GLSL linking failed!\ninfo log:\n{}\ndebug log:\n{}",
                  mHandle->getInfoLog(), mHandle->getInfoDebugLog()
        );
    }

    glslang::TIntermediate* referenceIntermediate = mHandle->getIntermediate(typeToGlslang(args.shaders[0]->getType()));
    SemanticIoResolver ioResolver(*referenceIntermediate, args.semanticMappings);
    if (!mHandle->mapIO(&ioResolver))
    {
        logAndDie("GLSL io mapping failed!\ninfo log:\n{}\ndebug log:\n{}",
                  mHandle->getInfoLog(), mHandle->getInfoDebugLog()
        );
    }

    mMeta = reflectProgram(*mHandle);
}

std::vector<std::uint32_t> GLSLShaderProgram::generateSpirv(vk::ShaderStageFlagBits stage) const
{
    const EShLanguage glslLang = typeToGlslang(stage);

    glslang::SpvOptions options = {};
    if (mLinkFlags.withDebugInfo)
    {
        options.generateDebugInfo = true;
        options.stripDebugInfo = false;
        options.disableOptimizer = true;
        options.emitNonSemanticShaderDebugInfo = true;
        options.emitNonSemanticShaderDebugSource = false; //  TODO: this should be true, but makes GLSLang crash
    }
    else
    {
        options.generateDebugInfo = false;
        options.stripDebugInfo = true;
        options.disableOptimizer = false;
        options.emitNonSemanticShaderDebugInfo = true; // TODO: this should be false, but that also crashes GLSLang ...
        options.emitNonSemanticShaderDebugSource = false;
    }
    options.optimizeSize = false;
    options.disassemble = false;
    options.validate = true;

    spv::SpvBuildLogger logger;
    const glslang::TIntermediate* intermediate = mHandle->getIntermediate(glslLang);
    if (intermediate == nullptr)
    {
        throw std::runtime_error("Attempting to generate SpirV from invalid shader stage.");
    }
    std::vector<std::uint32_t> spirv;
    glslang::GlslangToSpv(*intermediate, spirv, &logger, &options);

    const std::string messages = logger.getAllMessages();
    if (!messages.empty())
    {
        logMsg("SpirV messages: {}", messages);
    }

    return spirv;
}

std::vector<PipelineStage> GLSLShaderProgram::generatePipelineStages() const
{
    std::vector<PipelineStage> stages;
    for (const vk::ShaderStageFlagBits stage : mMeta.stages)
    {
        const std::vector<std::uint32_t> spirv = generateSpirv(stage);
        stages.push_back({
            .shader = getOwner()->createChild<ShaderModule>(ShaderModuleCreationArgs{.code = spirv}),
            .stage = stage
        });
    }
    return stages;
}

GraphicsPipelineCreationArgs GLSLShaderProgram::prepareGraphicsPipeline(PrepareGraphicsPipelineArgs& args) const
{
    args.pipelineLayoutMeta = mMeta.generatePipelineLayout(args.layoutArgs);
    args.layouts = args.pipelineLayoutMeta.createPipelineLayout(*getOwner());
    return {
        .stages = generatePipelineStages(),
        .vertexInput = mMeta.generateVertexInputFromLayout(args.vertexLayout),
        .layout = args.layouts.pipelineLayout
    };
}

ComputePipelineCreationArgs GLSLShaderProgram::prepareComputePipeline(PrepareComputePipelineArgs& args) const
{
    args.pipelineLayoutMeta = mMeta.generatePipelineLayout(args.layoutArgs);
    args.layouts = args.pipelineLayoutMeta.createPipelineLayout(*getOwner());
    return {
        .stage = generatePipelineStages()[0],
        .layout = args.layouts.pipelineLayout
    };
}

// GraphicsPipelineCreationArgs prepareGLSLGraphicsPipeline(const PrepareGLSLGraphicsPipelineArgs& args)
// {
//     return {
//         .stages =
//             {
//                 PipelineStage{.shader = vertexShaderModule, .stage = vk::ShaderStageFlagBits::eVertex},
//                 PipelineStage{.shader = fragmentShaderModule, .stage = vk::ShaderStageFlagBits::eFragment}
//             },
//         .vertexInput = std::move(vertexInput),
//         .inputAssembly =
//             {
//                 .topology = vk::PrimitiveTopology::eTriangleStrip
//             },
//         .colorBlend =
//             {
//                 .attachements = {DEFAULT_BLEND_ATTACHMENT}
//             },
//         .renderingInfo = GraphicsPipelineRenderingInfo{
//             .colorAttachmentFormats = {mRenderTarget->getFormat()}
//         },
//         .layout = mPipelineLayout
//     };
// }
} // namespace iwa