#define SPV_MAX_INTERFACE_VARIABLES 128
#define SPV_MAX_SEMANTIC_LENGTH 64

enum
{
    SpvMagic = 0x07230203,
    SpvOpExtension = 10,
    SpvOpEntryPoint = 15,
    SpvOpVariable = 59,
    SpvOpDecorate = 71,
    SpvOpDecorateStringGoogle = 5632,
    SpvOpMemberDecorateStringGoogle = 5633,
    SpvExecutionModelVertex = 0,
    SpvExecutionModelFragment = 4,
    SpvStorageClassInput = 1,
    SpvStorageClassOutput = 3,
    SpvDecorationBuiltIn = 11,
    SpvDecorationLocation = 30,
    SpvDecorationUserSemanticGoogle = 5635
};

struct spv_interface_variable
{
    uint32_t Id;
    uint32_t StorageClass;
    uint32_t *Location;
    char Semantic[SPV_MAX_SEMANTIC_LENGTH];
    bool IsBuiltIn;
};

struct spv_shader
{
    uint32_t *Words;
    uint64_t WordCount;
    rhi_shader_stage Stage;
    spv_interface_variable Variables[SPV_MAX_INTERFACE_VARIABLES];
    uint64_t VariableCount;
};

static uint32_t
FindLocationBySemantic(spv_shader* Shader, byteview semantic, uint32_t StorageClass)
{
    ASSERT(semantic.size <= COUNT(((spv_interface_variable*)0)->Semantic));
    
    for (uint32_t i = 0; i < Shader->VariableCount; ++i) {
        spv_interface_variable* var = &Shader->Variables[i];
        if (var->StorageClass == StorageClass &&
            !strncmp((char*)semantic.data, var->Semantic, semantic.size)) {
            return *var->Location;
        }
    }

    ASSERT(false, "%.*s", semantic.size, semantic.data);
}

static spv_interface_variable*
FindOrAddSpvInterfaceVariable(spv_shader* Shader, uint32_t Id)
{
    for (uint64_t i = 0; i < Shader->VariableCount; ++i) {
        if (Shader->Variables[i].Id == Id)
            return &Shader->Variables[i];
    }

    ASSERT(Shader->VariableCount < SPV_MAX_INTERFACE_VARIABLES);

    spv_interface_variable *ret = &Shader->Variables[Shader->VariableCount++];
    MEM_ZERO(ret, 0);
    ret->Id = Id;
    return ret;
}

static void
ProcessSpvShader(spv_shader *Shader, const uint32_t *in, uint64_t inCount, uint32_t *out)
{
    ASSERT(inCount >= 5);
    ASSERT(in[0] == SpvMagic);

    MEM_ZERO(Shader, 0);
    memcpy(out, in, 5 * sizeof(uint32_t));

    uint64_t inAt = 5;
    uint64_t outAt = 5;

    while (inAt < inCount) {
        uint32_t Instruction = in[inAt];
        uint32_t Op = Instruction & 0xffff;
        uint32_t WordCount = Instruction >> 16;
        bool KeepInstruction = true;

        ASSERT(WordCount != 0);
        ASSERT(inAt + WordCount <= inCount);

        switch (Op) {
        case SpvOpEntryPoint: {
            uint32_t ExecutionModel = in[inAt + 1];

            if (ExecutionModel == SpvExecutionModelVertex)
                Shader->Stage = Shader_Vertex;
            else if (ExecutionModel == SpvExecutionModelFragment)
                Shader->Stage = Shader_Pixel;
        } break;

        case SpvOpExtension: {
            const char *name = (const char *)&in[inAt + 1];

            if (!strcmp(name, "SPV_GOOGLE_decorate_string") ||
                !strcmp(name, "SPV_GOOGLE_hlsl_functionality1") ||
                !strcmp(name, "SPV_GOOGLE_user_type"))
                KeepInstruction = false;
        } break;

        case SpvOpDecorate: {
            uint32_t Id = in[inAt + 1];
            uint32_t Decoration = in[inAt + 2];

            if (Decoration == SpvDecorationLocation) {
                spv_interface_variable *Variable = FindOrAddSpvInterfaceVariable(Shader, Id);

                Variable->Location = &out[outAt + 3];
            }
            else if (Decoration == SpvDecorationBuiltIn) {
                spv_interface_variable *Variable = FindOrAddSpvInterfaceVariable(Shader, Id);

                Variable->IsBuiltIn = true;
            }
        } break;

        case SpvOpDecorateStringGoogle: {
            uint32_t Id = in[inAt + 1];
            uint32_t Decoration = in[inAt + 2];

            if (Decoration == SpvDecorationUserSemanticGoogle) {
                spv_interface_variable *Variable = FindOrAddSpvInterfaceVariable(Shader, Id);
                const char *Semantic = (const char *)&in[inAt + 3];
                uint64_t SemanticLength = (uint64_t)strlen(Semantic);

                ASSERT(SemanticLength < SPV_MAX_SEMANTIC_LENGTH);

                memcpy(Variable->Semantic, Semantic, SemanticLength + 1);
            }

            KeepInstruction = false;
        } break;

        case SpvOpMemberDecorateStringGoogle: {
            KeepInstruction = false;
        } break;

        case SpvOpVariable: {
            uint32_t Id = in[inAt + 2];
            uint32_t StorageClass = in[inAt + 3];

            if (StorageClass == SpvStorageClassInput ||
                StorageClass == SpvStorageClassOutput)
                {
                    spv_interface_variable *Variable =
                        FindOrAddSpvInterfaceVariable(Shader, Id);

                    Variable->StorageClass = StorageClass;
                }
        } break;
            
        }

        if (KeepInstruction) {
            memcpy(&out[outAt], &in[inAt], WordCount * sizeof(uint32_t));
            outAt += WordCount;
        }

        inAt += WordCount;
    }

    Shader->Words = out;
    Shader->WordCount = outAt;
}

static void
LinkSpvShaders(spv_shader *left, spv_shader *right)
{
    ASSERT(left->Stage  == Shader_Vertex);
    ASSERT(right->Stage == Shader_Pixel);

    for (uint64_t i = 0; i < right->VariableCount; ++i) {
        spv_interface_variable *varRight = &right->Variables[i];

        if (varRight->IsBuiltIn || varRight->Semantic[0] == '\0')
            continue;
        if (varRight->StorageClass != SpvStorageClassInput)
            continue;

        ASSERT(varRight->Location);

        bool linked = false;

        for (uint64_t j = 0; j < left->VariableCount; ++j) {
            spv_interface_variable *varLeft = &left->Variables[j];

            if (varLeft->IsBuiltIn || varLeft->Semantic[0] == '\0')
                continue;
            if (varLeft->StorageClass != SpvStorageClassOutput)
                continue;

            ASSERT(varLeft->Location);

            if (strncmp(varLeft->Semantic, varRight->Semantic, COUNT(varLeft->Semantic)))
                continue;

            *varRight->Location = *varLeft->Location;

            linked = true;
            break;
        }

        ASSERT(linked);
    }
}