#include "description_parser.h"
#include "processor.h"
#include "shader_compiler.h"

#include "runtime/buffer_manager.h"
#include "runtime/config.h"
#include "runtime/mem_arena.h"
#include "runtime/runtime.h"

#include "gfx/effect.h"
#include "gfx/gfx.h"

#include <assert.h>
#include <limits.h>
#include <stdbool.h>
#include <string.h>

typedef struct {
    rt_resource shaders[3];
    char *shader_names[3];
    unsigned int shader_count;

    unsigned int vertex_shader;
    unsigned int fragment_shader;
    unsigned int compute_shader;

    /* TODO Fixed function settings */

    /* Sampler settings */

    uint16_t uniform_binding_count;
    uint16_t storage_binding_count;
    uint16_t texture_binding_count;
} rt_parsed_pipeline_data;

typedef struct {
    unsigned int pass_count;
    rt_parsed_pipeline_data pipelines[RT_MAX_SUBRESOURCES];
    rt_render_pass_id pass_ids[RT_MAX_SUBRESOURCES];
} rt_parsed_effect_data;

enum {
    RT_SHADER_NOT_PRESENT = RT_ASSET_PROCESSING_FAILED + 1
};

extern RT_DLLIMPORT rt_cvar rt_Renderer;

static char *GenerateShaderName(rt_shader_type type,
                                rt_shader_stage stage,
                                rt_shader_optimization_level optimization,
                                const char *file_name,
                                rt_arena *arena) {
    size_t name_len = strlen(file_name) + 5 /* type */
                      + 5                   /* stage */
                      + 3                   /* optimization */
                      + 1;                  /* '\0' */
    char *res_name = rtArenaPush(arena, name_len);
    if (!res_name)
        return NULL;
    const char *type_str = NULL;
    switch (type) {
    case RT_SHADER_TYPE_VULKAN:
        type_str = ":vk";
        break;
    default:
        return NULL;
    }
    const char *stage_str = NULL;
    switch (stage) {
    case RT_SHADER_STAGE_VERTEX:
        stage_str = ":vert";
        break;
    case RT_SHADER_STAGE_FRAGMENT:
        stage_str = ":frag";
        break;
    case RT_SHADER_STAGE_COMPUTE:
        stage_str = ":comp";
        break;
    default:
        return NULL;
    }
    const char *optim_str = NULL;
    switch (optimization) {
    case RT_SHADER_OPTIMIZATION_NONE:
        optim_str = ":O0";
        break;
    case RT_SHADER_OPTIMIZATION_SIZE:
        optim_str = ":Os";
        break;
    case RT_SHADER_OPTIMIZATION_SPEED:
        optim_str = ":Ox";
        break;
    default:
        return NULL;
    }
    rtSPrint(res_name, name_len, "%s%s%s%s", file_name, type_str, stage_str, optim_str);
    return res_name;
}

static rt_result ParseShader(rt_parse_state *state,
                             unsigned int root_list,
                             const char *name,
                             const char *file_path,
                             rt_shader_type type,
                             rt_shader_stage stage,
                             rt_shader_optimization_level optimization,
                             rt_resource *p_resource,
                             char **p_resource_name,
                             rt_arena *arena) {
    const rt_parsed_stmt *stmt = rtFindStatement(state, root_list, name);
    if (stmt) {
        if (stmt->form != RT_STMT_FORM_LIST) {
            rtReportError("GFX",
                          "Expected a list as the value of "
                          "\"%s\" in %s",
                          name,
                          file_path);
            return RT_ASSET_PROCESSING_FAILED;
        }
        const rt_parsed_stmt_list *shader_list = &state->statement_lists[stmt->list_index];

        unsigned int stmt_index = shader_list->first;
        for (unsigned int i = 0; i < shader_list->count; ++i) {
            const rt_parsed_stmt *shader = &state->statements[stmt_index];
            if (shader->form != RT_STMT_FORM_VALUE && shader->form != RT_STMT_FORM_BLOCK) {
                rtReportError("GFX",
                              "Expected a simple statement or block for "
                              "\"%s.%*s\" in %s",
                              name,
                              (int)shader->attribute.length,
                              shader->attribute.start,
                              file_path);
                return RT_ASSET_PROCESSING_FAILED;
            }
            rt_shader_type in_file_type = RT_SHADER_TYPE_INVALID;
            if (rtCompareSpanToString(shader->attribute, "vk") == 0) {
                in_file_type = RT_SHADER_TYPE_VULKAN;
            } else {
                rtReportError("GFX",
                              "Invalid renderer backend"
                              "\"%*s\" in %s of file %s",
                              (int)shader->attribute.length,
                              shader->attribute.start,
                              name,
                              file_path);
                return RT_ASSET_PROCESSING_FAILED;
            }

            if (in_file_type == type) {
                if (shader->form == RT_STMT_FORM_BLOCK) {
                    /* Inline code */
                    rt_shader_bytecode bytecode =
                        CompileShader(type, stage, optimization, shader->block, file_path, arena);
                    if (!bytecode.bytes)
                        return RT_ASSET_PROCESSING_FAILED;

                    *p_resource_name =
                        GenerateShaderName(type, stage, optimization, file_path, arena);
                    if (!*p_resource_name)
                        return RT_ASSET_PROCESSING_FAILED;

                    rt_resource resource;
                    resource.type              = RT_RESOURCE_SHADER;
                    resource.dependency_count  = 0;
                    resource.subresource_count = 0;
                    resource.data = rtArenaPush(arena, sizeof(rt_shader_info) + bytecode.len);
                    if (!resource.data)
                        return RT_ASSET_PROCESSING_FAILED;
                    rt_shader_info *shader_info  = resource.data;
                    uint8_t *shader_bytecode     = (uint8_t *)(shader_info + 1);
                    shader_info->stage           = stage;
                    shader_info->type            = type;
                    shader_info->bytecode_length = bytecode.len;
                    rtSetRelptr(&shader_info->bytecode, shader_bytecode);
                    memcpy(shader_bytecode, bytecode.bytes, bytecode.len);

                    memcpy(p_resource, &resource, sizeof(resource));
                    break;
                } else if (shader->form != RT_STMT_FORM_VALUE) {
                    /* A filename */
                    rtLog("AC", "Shader source files not implemented yet!");
                    return RT_ASSET_PROCESSING_FAILED;
                }
            }

            stmt_index = shader->next;
        }
        return RT_SUCCESS;
    }
    return RT_SHADER_NOT_PRESENT;
}

static rt_shader_optimization_level
ParseOptimizationLevel(rt_parse_state *state, unsigned int root_list, const char *file_path) {
    const rt_parsed_stmt *stmt = rtFindStatement(state, root_list, "optimization");
    rt_shader_optimization_level optimization_level = RT_SHADER_OPTIMIZATION_NONE;
    if (stmt) {
        if (stmt->form != RT_STMT_FORM_VALUE) {
            rtReportError("GFX",
                          "Expected a simple statement for"
                          "\"optimization\" in %s",
                          file_path);
            return RT_SHADER_OPTIMIZATION_NONE;
        }

        if (rtCompareSpanToString(stmt->value, "speed") == 0) {
            optimization_level = RT_SHADER_OPTIMIZATION_SPEED;
        } else if (rtCompareSpanToString(stmt->value, "size") == 0) {
            optimization_level = RT_SHADER_OPTIMIZATION_SIZE;
        } else if (rtCompareSpanToString(stmt->value, "none") == 0) {
            optimization_level = RT_SHADER_OPTIMIZATION_NONE;
        } else {
            rtReportError("GFX",
                          "Expected one of 'speed', 'size' and 'none' for \"optimization\" in %s",
                          file_path);
        }
    }

    return optimization_level;
}

static rt_result ParsePipeline(rt_parse_state *state,
                               unsigned int root_list,
                               const char *file_path,
                               rt_shader_optimization_level optimization,
                               rt_parsed_pipeline_data *pipeline,
                               rt_arena *arena) {
    rt_result result = RT_SUCCESS;

    rt_shader_type type = RT_SHADER_TYPE_INVALID;
    if (strcmp(rt_Renderer.s, "vk") == 0)
        type = RT_SHADER_TYPE_VULKAN;

    if (type == RT_SHADER_TYPE_INVALID) {
        result = RT_ASSET_PROCESSING_FAILED;
        rtLog("AC", "rt_Renderer (%s) could not be translated to a shader type.", rt_Renderer.s);
        goto out;
    }

    /* Process shader stages */
    result = ParseShader(state,
                         root_list,
                         "vertex",
                         file_path,
                         type,
                         RT_SHADER_STAGE_VERTEX,
                         optimization,
                         &pipeline->shaders[pipeline->shader_count],
                         &pipeline->shader_names[pipeline->shader_count],
                         arena);
    if (result == RT_SUCCESS) {
        pipeline->vertex_shader = pipeline->shader_count;
        ++pipeline->shader_count;
    } else if (result == RT_SHADER_NOT_PRESENT) {
        pipeline->vertex_shader = UINT_MAX;
    } else {
        goto out;
    }
    result = RT_SUCCESS;

    result = ParseShader(state,
                         root_list,
                         "fragment",
                         file_path,
                         type,
                         RT_SHADER_STAGE_FRAGMENT,
                         optimization,
                         &pipeline->shaders[pipeline->shader_count],
                         &pipeline->shader_names[pipeline->shader_count],
                         arena);
    if (result == RT_SUCCESS) {
        pipeline->fragment_shader = pipeline->shader_count;
        ++pipeline->shader_count;
    } else if (result == RT_SHADER_NOT_PRESENT) {
        pipeline->fragment_shader = UINT_MAX;
    } else {
        goto out;
    }
    result = RT_SUCCESS;

    result = ParseShader(state,
                         root_list,
                         "compute",
                         file_path,
                         type,
                         RT_SHADER_STAGE_COMPUTE,
                         optimization,
                         &pipeline->shaders[pipeline->shader_count],
                         &pipeline->shader_names[pipeline->shader_count],
                         arena);
    if (result == RT_SUCCESS) {
        pipeline->compute_shader = pipeline->shader_count;
        ++pipeline->shader_count;
    } else if (result == RT_SHADER_NOT_PRESENT) {
        pipeline->compute_shader = UINT_MAX;
    } else {
        goto out;
    }
    result = RT_SUCCESS;

out:
    return result;
}

static rt_result ParseEffect(rt_file_id fid,
                             const char *text,
                             size_t length,
                             rt_parsed_effect_data *effect,
                             rt_arena *arena) {
    /* This is the grammar for pipeline files:
     * <stmt-list> ::= <stmt>*
     * <stmt> ::= <attribute> ( ( <value> ';' ) | ( '{' <stmt-list> '}' ) )
     * <attribute> ::= [:alnum:]*
     * <value>:: = [:alnum:]* */
    const char *file_path = rtGetFilePath(fid);
    rt_parse_state state;
    unsigned int root_list;
    rt_result result = rtParseDescription(text, length, file_path, &root_list, &state, arena);
    if (result != RT_SUCCESS) {
        return result;
    }

    memset(effect, 0, sizeof(*effect));

    rt_shader_optimization_level optimization =
        ParseOptimizationLevel(&state, root_list, file_path);

    const rt_parsed_stmt *passes_stmt = rtFindStatement(&state, root_list, "passes");
    if (!passes_stmt) {
        rtLog("AC", "Did not find passes list in %s.", file_path);
        return RT_INVALID_VALUE;
    }
    const rt_parsed_stmt_list *passes_list = &state.statement_lists[passes_stmt->list_index];
    if (passes_list->count > RT_MAX_SUBRESOURCES) {
        rtLog("AC",
              "Too many passes in 'passes'. Maximum supported number is %u. In %s.",
              RT_MAX_SUBRESOURCES,
              file_path);
        return RT_INVALID_VALUE;
    }
    effect->pass_count              = passes_list->count;
    const rt_parsed_stmt *pass_stmt = &state.statements[passes_list->first];
    for (unsigned int i = 0; i < passes_list->count;
         ++i, pass_stmt = &state.statements[pass_stmt->next]) {
        if (pass_stmt->form != RT_STMT_FORM_LIST) {
            rtLog("AC",
                  "Expected a list as the value of passes.%*s in %s",
                  pass_stmt->attribute.length,
                  pass_stmt->attribute.start,
                  file_path);
            return RT_INVALID_VALUE;
        }
        effect->pass_ids[i] =
            rtCalculateRenderPassID(pass_stmt->attribute.start, pass_stmt->attribute.length);

        result = ParsePipeline(&state,
                               pass_stmt->list_index,
                               file_path,
                               optimization,
                               &effect->pipelines[i],
                               arena);
        if (result != RT_SUCCESS)
            return result;
    }

    return result;
}

RT_ASSET_PROCESSOR_FN(EffectProcessor) {
    rt_loaded_asset asset = LoadAsset(file);
    if (!asset.buffer)
        return RT_UNKNOWN_ERROR;

    rt_parsed_effect_data effect;
    memset(&effect, 0, sizeof(effect));
    rt_result result = ParseEffect(file, asset.buffer, asset.size, &effect, arena);
    if (result != RT_SUCCESS)
        goto out;

    rt_effect_info effect_info;
    effect_info.pass_count = effect.pass_count;

    rt_resource effect_resource       = {0};
    effect_resource.data              = &effect_info;
    effect_resource.type              = RT_RESOURCE_EFFECT;
    effect_resource.subresource_count = effect.pass_count;

    *new_resource_count = 0;

    const char *name = rtGetFilePath(file);
    for (unsigned int i = 0; i < effect.pass_count; ++i) {
        effect_info.passes[i].pass_id  = effect.pass_ids[i];
        effect_info.passes[i].pipeline = RT_INVALID_RESOURCE_ID;

        rt_parsed_pipeline_data pipeline;
        memcpy(&pipeline, &effect.pipelines[i], sizeof(pipeline));

        rt_resource_id shader_resources[3] = {0};
        result                             = rtCreateResources(pipeline.shader_count,
                                   (const char **)pipeline.shader_names,
                                   pipeline.shaders,
                                   shader_resources);
        if (result != RT_SUCCESS)
            goto out;

        rt_resource pipeline_resource      = {0};
        pipeline_resource.type             = RT_RESOURCE_PIPELINE;
        pipeline_resource.dependency_count = pipeline.shader_count;
        memcpy(pipeline_resource.dependencies, shader_resources, sizeof(shader_resources));
        pipeline_resource.subresource_count = 0;

        size_t data_size       = sizeof(rt_pipeline_info);
        pipeline_resource.data = rtArenaPush(arena, data_size);
        if (!pipeline_resource.data) {
            result = RT_OUT_OF_MEMORY;
            goto out;
        }
        rt_pipeline_info *info = pipeline_resource.data;
        memset(info, 0, sizeof(*info));
        info->vertex_shader   = (pipeline.vertex_shader != UINT_MAX)
                                    ? shader_resources[pipeline.vertex_shader]
                                    : RT_INVALID_RESOURCE_ID;
        info->fragment_shader = (pipeline.fragment_shader != UINT_MAX)
                                    ? shader_resources[pipeline.fragment_shader]
                                    : RT_INVALID_RESOURCE_ID;
        info->compute_shader  = (pipeline.compute_shader != UINT_MAX)
                                    ? shader_resources[pipeline.compute_shader]
                                    : RT_INVALID_RESOURCE_ID;
        rt_resource_id pipeline_id;
        char pipeline_name[260];
        rtSPrint(pipeline_name, sizeof(pipeline_name), "%s:%u", name, i);

        char *ppln = &pipeline_name[0];

        result = rtCreateResources(1, &ppln, &pipeline_resource, &pipeline_id);
        if (result == RT_SUCCESS) {
            new_resources[i] = pipeline_id;
            memcpy(&new_resources[i + 1], shader_resources, sizeof(shader_resources));
            *new_resource_count += 1 + pipeline.shader_count;
            effect_resource.subresources[i] = pipeline_id;
            effect_info.passes[i].pipeline  = pipeline_id;
        }
    }

    rt_resource_id effect_id = 0;
    result                   = rtCreateResources(1, &name, &effect_resource, &effect_id);
    if (result == RT_SUCCESS) {
        new_resources[*new_resource_count] = effect_id;
        *new_resource_count += 1;
    }
out:
    rtLog("AC", "Released %p", asset.buffer);
    rtReleaseBuffer(asset.buffer, asset.size);
    return result;
}