// SPDX-License-Identifier: MIT
// mvgal_spirv_opt.cpp — SPIRV-Tools optimizer wrapper (Phase 2)

#include "mvgal_shader_translate.h"

#include <cstdlib>
#include <cstring>
#include <new>
#include <vector>

#include "spirv-tools/optimizer.hpp"

struct spirv_opt_ctx {
    spvtools::Optimizer *opt;
    spvtools::OptimizerOptions opts;
};

extern "C" void *mvgal_spirv_opt_create(void) {
    auto ctx = new (std::nothrow) spirv_opt_ctx{};
    if (!ctx) return nullptr;

    ctx->opt = new (std::nothrow) spvtools::Optimizer(SPV_ENV_VULKAN_1_3);
    if (!ctx->opt) { delete ctx; return nullptr; }

    // Register standard optimization passes (SPIRV-Tools v2024.2 API)
    ctx->opt->RegisterPass(spvtools::CreateFreezeSpecConstantValuePass());
    ctx->opt->RegisterPass(spvtools::CreateFoldSpecConstantOpAndCompositePass());
    ctx->opt->RegisterPass(spvtools::CreateEliminateDeadFunctionsPass());
    ctx->opt->RegisterPass(spvtools::CreateEliminateDeadMembersPass());
    ctx->opt->RegisterPass(spvtools::CreateInlineExhaustivePass());
    ctx->opt->RegisterPass(spvtools::CreateLocalAccessChainConvertPass());
    ctx->opt->RegisterPass(spvtools::CreateLocalSingleBlockLoadStoreElimPass());
    ctx->opt->RegisterPass(spvtools::CreateLocalSingleStoreElimPass());
    ctx->opt->RegisterPass(spvtools::CreateSimplificationPass());
    ctx->opt->RegisterPass(spvtools::CreateVectorDCEPass());
    ctx->opt->RegisterPass(spvtools::CreateBlockMergePass());

    ctx->opts.set_run_validator(false); // assume valid input from Vulkan driver

    return ctx;
}

extern "C" void mvgal_spirv_opt_destroy(void *ctx) {
    if (!ctx) return;
    auto c = static_cast<spirv_opt_ctx *>(ctx);
    delete c->opt;
    delete c;
}

extern "C" int mvgal_spirv_optimize(void *ctx,
                                    const uint32_t *spirv,
                                    size_t word_count,
                                    uint32_t **out_spirv,
                                    size_t *out_word_count) {
    if (!ctx || !spirv || !out_spirv || !out_word_count)
        return -1;

    auto c = static_cast<spirv_opt_ctx *>(ctx);
    std::vector<uint32_t> optimized;

    if (!c->opt->Run(spirv, word_count, &optimized, c->opts))
        return -1;

    *out_word_count = optimized.size();
    *out_spirv = static_cast<uint32_t *>(std::malloc(optimized.size() * sizeof(uint32_t)));
    if (!*out_spirv) return -1;

    std::memcpy(*out_spirv, optimized.data(), optimized.size() * sizeof(uint32_t));
    return 0;
}
