#include "asset_compiler.h"
#include "config.h"
#include "shader_compiler.h"
#include "buffer_manager.h"

#include <shaderc/shaderc.h>
#include <string.h>

RT_CVAR_S(rt_VkShaderCompilerShaderBasepath,
          "Basepath for standard includes (<source>). Default: assets/shader/",
          "assets/shader/");

static shaderc_include_result *IncludeResolveCallback(void *user_data,
                                                      const char *requested_source,
                                                      int type,
                                                      const char *requesting_source,
                                                      size_t include_depth) {
    RT_UNUSED(include_depth);
    rt_arena *arena = user_data;

    shaderc_include_result *result = RT_ARENA_PUSH_STRUCT_ZERO(arena, shaderc_include_result);
    if (!result)
        return NULL;

    /* Resolve the path */
    const char *basepath = NULL;
    switch (type) {
    case shaderc_include_type_standard:
        basepath = rt_VkShaderCompilerShaderBasepath.s;
        break;
    case shaderc_include_type_relative: {
        size_t len = strlen(requesting_source);
        char *tmp  = rtArenaPush(arena, len + 1);
        if (!tmp) {
            rtReportError("AC", "Out of memory in the vulkan shader compiler.");
            result->content = "Out of memory in the vulkan shader compiler.";
            result->content_length = strlen(result->content);
            return result;
        }
        memcpy(tmp, requesting_source, len + 1);
        while (len > 0 && tmp[len] != '/' && tmp[len] != '\\')
            --len;
        tmp[len + 1] = '\0';
        basepath     = tmp;
    } break;
    default:
        rtReportError("AC", "Invalid include type reported in %s", requesting_source);
        return result;
    }
    size_t basepath_len  = strlen(basepath);
    size_t requested_len = strlen(requested_source);
    char *full_path      = rtArenaPush(arena, basepath_len + requested_len + 1);
    if (!full_path) {
        rtReportError("AC", "Out of memory in vulkan shader compiler.");
        result->content        = "Out of memory in the vulkan shader compiler.";
        result->content_length = strlen(result->content);
    }
    memcpy(full_path, basepath, basepath_len);
    memcpy(full_path + basepath_len, requested_source, requested_len);
    full_path[basepath_len + requested_len] = '\0';

    rt_loaded_asset source = LoadAsset(rtAddFile(full_path));
    if (source.buffer) {
        result->content        = source.buffer;
        result->content_length = source.size;
        result->source_name    = full_path;
        result->source_name_length = basepath_len + requested_len + 1;
    } else {
        char *error_tmpl = "Could not open ";
        char *error_msg = rtArenaPush(arena, strlen(error_tmpl) + basepath_len + requested_len + 1);
        if (error_msg) {
            memcpy(error_msg, error_tmpl, strlen(error_tmpl));
            memcpy(error_msg + strlen(error_tmpl), full_path, basepath_len + requested_len + 1);
            result->content = error_msg;
            result->content_length = strlen(error_tmpl) + basepath_len + requested_len;
        } else {
            result->content        = "Could not open header file.";
            result->content_length = strlen(result->content);
        }
    }

    return result;
}

static void IncludeReleaseCallback(void *user_data, shaderc_include_result *include_result) {
    RT_UNUSED(user_data);
    if (include_result->source_name) {
        /* Include result contains a valid file buffer */
        rtReleaseBuffer(include_result->content, include_result->content_length);
    }
}

rt_shader_bytecode CompileVulkanShader(rt_shader_stage stage,
                                       rt_shader_optimization_level optimization,
                                       rt_text_span code,
                                       const char *file_path,
                                       rt_arena *arena) {
    rt_arena_rewindpoint arena_rewind = rtGetArenaRewindPoint(arena);

    shaderc_shader_kind kind;
    const char *entry;
    switch (stage) {
    case RT_SHADER_STAGE_VERTEX:
        kind  = shaderc_vertex_shader;
        entry = "VsMain";
        break;
    case RT_SHADER_STAGE_FRAGMENT:
        kind  = shaderc_fragment_shader;
        entry = "PsMain";
        break;
    case RT_SHADER_STAGE_COMPUTE:
        kind  = shaderc_compute_shader;
        entry = "CsMain";
        break;
    default:
        rtLog("AC", "Invalid shader stage %u", stage);
        return (rt_shader_bytecode){.bytes = NULL, .len = 0};
    }

    shaderc_compiler_t compiler = shaderc_compiler_initialize();
    if (!compiler) {
        rtLog("AC", "Failed to initialize the shaderc compiler.");
        return (rt_shader_bytecode){.bytes = NULL, .len = 0};
    }
    shaderc_compile_options_t options = shaderc_compile_options_initialize();
    if (!options) {
        rtLog("AC", "Failed to initialize the shaderc compile options.");
        shaderc_compiler_release(compiler);
        return (rt_shader_bytecode){.bytes = NULL, .len = 0};
    }
    shaderc_compile_options_add_macro_definition(options, "RT_VULKAN", 9, NULL, 0);
    shaderc_compile_options_set_source_language(options, shaderc_source_language_hlsl);
    if (optimization == RT_SHADER_OPTIMIZATION_SIZE)
        shaderc_compile_options_set_optimization_level(options, shaderc_optimization_level_size);
    else if (optimization == RT_SHADER_OPTIMIZATION_SPEED)
        shaderc_compile_options_set_optimization_level(options,
                                                       shaderc_optimization_level_performance);
    shaderc_compile_options_set_include_callbacks(options,
                                                  IncludeResolveCallback,
                                                  IncludeReleaseCallback,
                                                  arena);

    shaderc_compilation_result_t res = shaderc_compile_into_spv(compiler,
                                                                code.start,
                                                                code.length,
                                                                kind,
                                                                file_path,
                                                                entry,
                                                                options);
    rt_shader_bytecode bytecode      = {.bytes = NULL, .len = 0};
    if (shaderc_result_get_compilation_status(res) != shaderc_compilation_status_success) {
        rtLog("AC", "Shader compilation failed: %s", shaderc_result_get_error_message(res));
        goto out;
    }
    bytecode.len = shaderc_result_get_length(res);
    if (bytecode.len > 0) {
        bytecode.bytes = rtArenaPush(arena, bytecode.len);
        if (!bytecode.bytes) {
            rtLog("AC", "Failed to allocate %zu bytes for shader spv.", bytecode.len);
            bytecode.len = 0;
            goto out;
        }
        memcpy(bytecode.bytes, shaderc_result_get_bytes(res), bytecode.len);
    }

out:
    shaderc_compile_options_release(options);
    shaderc_compiler_release(compiler);
    rtArenaRewind(arena, arena_rewind);
    return bytecode;
}