#include <android/log.h>
#include <jni.h>

#define VULKAN_HPP_NO_STRUCT_CONSTRUCTORS

#include <mbgl/style/layers/custom_layer_host.hpp>
#include <mbgl/style/layers/vulkan/custom_layer_init_parameters.hpp>
#include <mbgl/style/layers/vulkan/custom_layer_render_parameters.hpp>

#include <array>
#include <vector>

const char *VK_LOG_TAG = "VulkanCustomLayer";

// ---------------------------------------------------------------------------
// Minimal SPIR-V shaders
//
// Vertex shader: generates a full-screen quad from gl_VertexIndex (no VBO).
//   index 0 → (-1,-1)  index 1 → ( 3,-1)  index 2 → (-1, 3)
//   This covers the NDC clip space with a single triangle.
//
// Fragment shader: outputs a push-constant RGBA color.
//
// Generated from:
//   vert: #version 450
//         void main() {
//           vec2 pos[3] = vec2[](vec2(-1,-1), vec2(3,-1), vec2(-1,3));
//           gl_Position = vec4(pos[gl_VertexIndex], 0.0, 1.0);
//         }
//   frag: #version 450
//         layout(push_constant) uniform PC { vec4 color; } pc;
//         layout(location = 0) out vec4 outColor;
//         void main() { outColor = pc.color; }
// ---------------------------------------------------------------------------

// clang-format off
static const unsigned char kVertSpirv[] = {
        0x03, 0x02, 0x23, 0x07, 0x00, 0x00, 0x01, 0x00, 0x0b, 0x00, 0x0d, 0x00,
        0x28, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x11, 0x00, 0x02, 0x00,
        0x01, 0x00, 0x00, 0x00, 0x0b, 0x00, 0x06, 0x00, 0x01, 0x00, 0x00, 0x00,
        0x47, 0x4c, 0x53, 0x4c, 0x2e, 0x73, 0x74, 0x64, 0x2e, 0x34, 0x35, 0x30,
        0x00, 0x00, 0x00, 0x00, 0x0e, 0x00, 0x03, 0x00, 0x00, 0x00, 0x00, 0x00,
        0x01, 0x00, 0x00, 0x00, 0x0f, 0x00, 0x07, 0x00, 0x00, 0x00, 0x00, 0x00,
        0x04, 0x00, 0x00, 0x00, 0x6d, 0x61, 0x69, 0x6e, 0x00, 0x00, 0x00, 0x00,
        0x18, 0x00, 0x00, 0x00, 0x1c, 0x00, 0x00, 0x00, 0x03, 0x00, 0x03, 0x00,
        0x02, 0x00, 0x00, 0x00, 0xc2, 0x01, 0x00, 0x00, 0x04, 0x00, 0x0a, 0x00,
        0x47, 0x4c, 0x5f, 0x47, 0x4f, 0x4f, 0x47, 0x4c, 0x45, 0x5f, 0x63, 0x70,
        0x70, 0x5f, 0x73, 0x74, 0x79, 0x6c, 0x65, 0x5f, 0x6c, 0x69, 0x6e, 0x65,
        0x5f, 0x64, 0x69, 0x72, 0x65, 0x63, 0x74, 0x69, 0x76, 0x65, 0x00, 0x00,
        0x04, 0x00, 0x08, 0x00, 0x47, 0x4c, 0x5f, 0x47, 0x4f, 0x4f, 0x47, 0x4c,
        0x45, 0x5f, 0x69, 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, 0x5f, 0x64, 0x69,
        0x72, 0x65, 0x63, 0x74, 0x69, 0x76, 0x65, 0x00, 0x05, 0x00, 0x04, 0x00,
        0x04, 0x00, 0x00, 0x00, 0x6d, 0x61, 0x69, 0x6e, 0x00, 0x00, 0x00, 0x00,
        0x05, 0x00, 0x03, 0x00, 0x0c, 0x00, 0x00, 0x00, 0x70, 0x6f, 0x73, 0x00,
        0x05, 0x00, 0x06, 0x00, 0x16, 0x00, 0x00, 0x00, 0x67, 0x6c, 0x5f, 0x50,
        0x65, 0x72, 0x56, 0x65, 0x72, 0x74, 0x65, 0x78, 0x00, 0x00, 0x00, 0x00,
        0x06, 0x00, 0x06, 0x00, 0x16, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
        0x67, 0x6c, 0x5f, 0x50, 0x6f, 0x73, 0x69, 0x74, 0x69, 0x6f, 0x6e, 0x00,
        0x06, 0x00, 0x07, 0x00, 0x16, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00,
        0x67, 0x6c, 0x5f, 0x50, 0x6f, 0x69, 0x6e, 0x74, 0x53, 0x69, 0x7a, 0x65,
        0x00, 0x00, 0x00, 0x00, 0x06, 0x00, 0x07, 0x00, 0x16, 0x00, 0x00, 0x00,
        0x02, 0x00, 0x00, 0x00, 0x67, 0x6c, 0x5f, 0x43, 0x6c, 0x69, 0x70, 0x44,
        0x69, 0x73, 0x74, 0x61, 0x6e, 0x63, 0x65, 0x00, 0x06, 0x00, 0x07, 0x00,
        0x16, 0x00, 0x00, 0x00, 0x03, 0x00, 0x00, 0x00, 0x67, 0x6c, 0x5f, 0x43,
        0x75, 0x6c, 0x6c, 0x44, 0x69, 0x73, 0x74, 0x61, 0x6e, 0x63, 0x65, 0x00,
        0x05, 0x00, 0x03, 0x00, 0x18, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
        0x05, 0x00, 0x06, 0x00, 0x1c, 0x00, 0x00, 0x00, 0x67, 0x6c, 0x5f, 0x56,
        0x65, 0x72, 0x74, 0x65, 0x78, 0x49, 0x6e, 0x64, 0x65, 0x78, 0x00, 0x00,
        0x47, 0x00, 0x03, 0x00, 0x16, 0x00, 0x00, 0x00, 0x02, 0x00, 0x00, 0x00,
        0x48, 0x00, 0x05, 0x00, 0x16, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
        0x0b, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x48, 0x00, 0x05, 0x00,
        0x16, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x0b, 0x00, 0x00, 0x00,
        0x01, 0x00, 0x00, 0x00, 0x48, 0x00, 0x05, 0x00, 0x16, 0x00, 0x00, 0x00,
        0x02, 0x00, 0x00, 0x00, 0x0b, 0x00, 0x00, 0x00, 0x03, 0x00, 0x00, 0x00,
        0x48, 0x00, 0x05, 0x00, 0x16, 0x00, 0x00, 0x00, 0x03, 0x00, 0x00, 0x00,
        0x0b, 0x00, 0x00, 0x00, 0x04, 0x00, 0x00, 0x00, 0x47, 0x00, 0x04, 0x00,
        0x1c, 0x00, 0x00, 0x00, 0x0b, 0x00, 0x00, 0x00, 0x2a, 0x00, 0x00, 0x00,
        0x13, 0x00, 0x02, 0x00, 0x02, 0x00, 0x00, 0x00, 0x21, 0x00, 0x03, 0x00,
        0x03, 0x00, 0x00, 0x00, 0x02, 0x00, 0x00, 0x00, 0x16, 0x00, 0x03, 0x00,
        0x06, 0x00, 0x00, 0x00, 0x20, 0x00, 0x00, 0x00, 0x17, 0x00, 0x04, 0x00,
        0x07, 0x00, 0x00, 0x00, 0x06, 0x00, 0x00, 0x00, 0x02, 0x00, 0x00, 0x00,
        0x15, 0x00, 0x04, 0x00, 0x08, 0x00, 0x00, 0x00, 0x20, 0x00, 0x00, 0x00,
        0x00, 0x00, 0x00, 0x00, 0x2b, 0x00, 0x04, 0x00, 0x08, 0x00, 0x00, 0x00,
        0x09, 0x00, 0x00, 0x00, 0x03, 0x00, 0x00, 0x00, 0x1c, 0x00, 0x04, 0x00,
        0x0a, 0x00, 0x00, 0x00, 0x07, 0x00, 0x00, 0x00, 0x09, 0x00, 0x00, 0x00,
        0x20, 0x00, 0x04, 0x00, 0x0b, 0x00, 0x00, 0x00, 0x07, 0x00, 0x00, 0x00,
        0x0a, 0x00, 0x00, 0x00, 0x2b, 0x00, 0x04, 0x00, 0x06, 0x00, 0x00, 0x00,
        0x0d, 0x00, 0x00, 0x00, 0x00, 0x00, 0x80, 0xbf, 0x2c, 0x00, 0x05, 0x00,
        0x07, 0x00, 0x00, 0x00, 0x0e, 0x00, 0x00, 0x00, 0x0d, 0x00, 0x00, 0x00,
        0x0d, 0x00, 0x00, 0x00, 0x2b, 0x00, 0x04, 0x00, 0x06, 0x00, 0x00, 0x00,
        0x0f, 0x00, 0x00, 0x00, 0x00, 0x00, 0x40, 0x40, 0x2c, 0x00, 0x05, 0x00,
        0x07, 0x00, 0x00, 0x00, 0x10, 0x00, 0x00, 0x00, 0x0f, 0x00, 0x00, 0x00,
        0x0d, 0x00, 0x00, 0x00, 0x2c, 0x00, 0x05, 0x00, 0x07, 0x00, 0x00, 0x00,
        0x11, 0x00, 0x00, 0x00, 0x0d, 0x00, 0x00, 0x00, 0x0f, 0x00, 0x00, 0x00,
        0x2c, 0x00, 0x06, 0x00, 0x0a, 0x00, 0x00, 0x00, 0x12, 0x00, 0x00, 0x00,
        0x0e, 0x00, 0x00, 0x00, 0x10, 0x00, 0x00, 0x00, 0x11, 0x00, 0x00, 0x00,
        0x17, 0x00, 0x04, 0x00, 0x13, 0x00, 0x00, 0x00, 0x06, 0x00, 0x00, 0x00,
        0x04, 0x00, 0x00, 0x00, 0x2b, 0x00, 0x04, 0x00, 0x08, 0x00, 0x00, 0x00,
        0x14, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x1c, 0x00, 0x04, 0x00,
        0x15, 0x00, 0x00, 0x00, 0x06, 0x00, 0x00, 0x00, 0x14, 0x00, 0x00, 0x00,
        0x1e, 0x00, 0x06, 0x00, 0x16, 0x00, 0x00, 0x00, 0x13, 0x00, 0x00, 0x00,
        0x06, 0x00, 0x00, 0x00, 0x15, 0x00, 0x00, 0x00, 0x15, 0x00, 0x00, 0x00,
        0x20, 0x00, 0x04, 0x00, 0x17, 0x00, 0x00, 0x00, 0x03, 0x00, 0x00, 0x00,
        0x16, 0x00, 0x00, 0x00, 0x3b, 0x00, 0x04, 0x00, 0x17, 0x00, 0x00, 0x00,
        0x18, 0x00, 0x00, 0x00, 0x03, 0x00, 0x00, 0x00, 0x15, 0x00, 0x04, 0x00,
        0x19, 0x00, 0x00, 0x00, 0x20, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00,
        0x2b, 0x00, 0x04, 0x00, 0x19, 0x00, 0x00, 0x00, 0x1a, 0x00, 0x00, 0x00,
        0x00, 0x00, 0x00, 0x00, 0x20, 0x00, 0x04, 0x00, 0x1b, 0x00, 0x00, 0x00,
        0x01, 0x00, 0x00, 0x00, 0x19, 0x00, 0x00, 0x00, 0x3b, 0x00, 0x04, 0x00,
        0x1b, 0x00, 0x00, 0x00, 0x1c, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00,
        0x20, 0x00, 0x04, 0x00, 0x1e, 0x00, 0x00, 0x00, 0x07, 0x00, 0x00, 0x00,
        0x07, 0x00, 0x00, 0x00, 0x2b, 0x00, 0x04, 0x00, 0x06, 0x00, 0x00, 0x00,
        0x21, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x2b, 0x00, 0x04, 0x00,
        0x06, 0x00, 0x00, 0x00, 0x22, 0x00, 0x00, 0x00, 0x00, 0x00, 0x80, 0x3f,
        0x20, 0x00, 0x04, 0x00, 0x26, 0x00, 0x00, 0x00, 0x03, 0x00, 0x00, 0x00,
        0x13, 0x00, 0x00, 0x00, 0x36, 0x00, 0x05, 0x00, 0x02, 0x00, 0x00, 0x00,
        0x04, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x03, 0x00, 0x00, 0x00,
        0xf8, 0x00, 0x02, 0x00, 0x05, 0x00, 0x00, 0x00, 0x3b, 0x00, 0x04, 0x00,
        0x0b, 0x00, 0x00, 0x00, 0x0c, 0x00, 0x00, 0x00, 0x07, 0x00, 0x00, 0x00,
        0x3e, 0x00, 0x03, 0x00, 0x0c, 0x00, 0x00, 0x00, 0x12, 0x00, 0x00, 0x00,
        0x3d, 0x00, 0x04, 0x00, 0x19, 0x00, 0x00, 0x00, 0x1d, 0x00, 0x00, 0x00,
        0x1c, 0x00, 0x00, 0x00, 0x41, 0x00, 0x05, 0x00, 0x1e, 0x00, 0x00, 0x00,
        0x1f, 0x00, 0x00, 0x00, 0x0c, 0x00, 0x00, 0x00, 0x1d, 0x00, 0x00, 0x00,
        0x3d, 0x00, 0x04, 0x00, 0x07, 0x00, 0x00, 0x00, 0x20, 0x00, 0x00, 0x00,
        0x1f, 0x00, 0x00, 0x00, 0x51, 0x00, 0x05, 0x00, 0x06, 0x00, 0x00, 0x00,
        0x23, 0x00, 0x00, 0x00, 0x20, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
        0x51, 0x00, 0x05, 0x00, 0x06, 0x00, 0x00, 0x00, 0x24, 0x00, 0x00, 0x00,
        0x20, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x50, 0x00, 0x07, 0x00,
        0x13, 0x00, 0x00, 0x00, 0x25, 0x00, 0x00, 0x00, 0x23, 0x00, 0x00, 0x00,
        0x24, 0x00, 0x00, 0x00, 0x21, 0x00, 0x00, 0x00, 0x22, 0x00, 0x00, 0x00,
        0x41, 0x00, 0x05, 0x00, 0x26, 0x00, 0x00, 0x00, 0x27, 0x00, 0x00, 0x00,
        0x18, 0x00, 0x00, 0x00, 0x1a, 0x00, 0x00, 0x00, 0x3e, 0x00, 0x03, 0x00,
        0x27, 0x00, 0x00, 0x00, 0x25, 0x00, 0x00, 0x00, 0xfd, 0x00, 0x01, 0x00,
        0x38, 0x00, 0x01, 0x00
};

static const unsigned char kFragSpirv[] = {
        0x03, 0x02, 0x23, 0x07, 0x00, 0x00, 0x01, 0x00, 0x0b, 0x00, 0x0d, 0x00,
        0x12, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x11, 0x00, 0x02, 0x00,
        0x01, 0x00, 0x00, 0x00, 0x0b, 0x00, 0x06, 0x00, 0x01, 0x00, 0x00, 0x00,
        0x47, 0x4c, 0x53, 0x4c, 0x2e, 0x73, 0x74, 0x64, 0x2e, 0x34, 0x35, 0x30,
        0x00, 0x00, 0x00, 0x00, 0x0e, 0x00, 0x03, 0x00, 0x00, 0x00, 0x00, 0x00,
        0x01, 0x00, 0x00, 0x00, 0x0f, 0x00, 0x06, 0x00, 0x04, 0x00, 0x00, 0x00,
        0x04, 0x00, 0x00, 0x00, 0x6d, 0x61, 0x69, 0x6e, 0x00, 0x00, 0x00, 0x00,
        0x09, 0x00, 0x00, 0x00, 0x10, 0x00, 0x03, 0x00, 0x04, 0x00, 0x00, 0x00,
        0x07, 0x00, 0x00, 0x00, 0x03, 0x00, 0x03, 0x00, 0x02, 0x00, 0x00, 0x00,
        0xc2, 0x01, 0x00, 0x00, 0x04, 0x00, 0x0a, 0x00, 0x47, 0x4c, 0x5f, 0x47,
        0x4f, 0x4f, 0x47, 0x4c, 0x45, 0x5f, 0x63, 0x70, 0x70, 0x5f, 0x73, 0x74,
        0x79, 0x6c, 0x65, 0x5f, 0x6c, 0x69, 0x6e, 0x65, 0x5f, 0x64, 0x69, 0x72,
        0x65, 0x63, 0x74, 0x69, 0x76, 0x65, 0x00, 0x00, 0x04, 0x00, 0x08, 0x00,
        0x47, 0x4c, 0x5f, 0x47, 0x4f, 0x4f, 0x47, 0x4c, 0x45, 0x5f, 0x69, 0x6e,
        0x63, 0x6c, 0x75, 0x64, 0x65, 0x5f, 0x64, 0x69, 0x72, 0x65, 0x63, 0x74,
        0x69, 0x76, 0x65, 0x00, 0x05, 0x00, 0x04, 0x00, 0x04, 0x00, 0x00, 0x00,
        0x6d, 0x61, 0x69, 0x6e, 0x00, 0x00, 0x00, 0x00, 0x05, 0x00, 0x05, 0x00,
        0x09, 0x00, 0x00, 0x00, 0x6f, 0x75, 0x74, 0x43, 0x6f, 0x6c, 0x6f, 0x72,
        0x00, 0x00, 0x00, 0x00, 0x05, 0x00, 0x03, 0x00, 0x0a, 0x00, 0x00, 0x00,
        0x50, 0x43, 0x00, 0x00, 0x06, 0x00, 0x05, 0x00, 0x0a, 0x00, 0x00, 0x00,
        0x00, 0x00, 0x00, 0x00, 0x63, 0x6f, 0x6c, 0x6f, 0x72, 0x00, 0x00, 0x00,
        0x05, 0x00, 0x03, 0x00, 0x0c, 0x00, 0x00, 0x00, 0x70, 0x63, 0x00, 0x00,
        0x47, 0x00, 0x04, 0x00, 0x09, 0x00, 0x00, 0x00, 0x1e, 0x00, 0x00, 0x00,
        0x00, 0x00, 0x00, 0x00, 0x47, 0x00, 0x03, 0x00, 0x0a, 0x00, 0x00, 0x00,
        0x02, 0x00, 0x00, 0x00, 0x48, 0x00, 0x05, 0x00, 0x0a, 0x00, 0x00, 0x00,
        0x00, 0x00, 0x00, 0x00, 0x23, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
        0x13, 0x00, 0x02, 0x00, 0x02, 0x00, 0x00, 0x00, 0x21, 0x00, 0x03, 0x00,
        0x03, 0x00, 0x00, 0x00, 0x02, 0x00, 0x00, 0x00, 0x16, 0x00, 0x03, 0x00,
        0x06, 0x00, 0x00, 0x00, 0x20, 0x00, 0x00, 0x00, 0x17, 0x00, 0x04, 0x00,
        0x07, 0x00, 0x00, 0x00, 0x06, 0x00, 0x00, 0x00, 0x04, 0x00, 0x00, 0x00,
        0x20, 0x00, 0x04, 0x00, 0x08, 0x00, 0x00, 0x00, 0x03, 0x00, 0x00, 0x00,
        0x07, 0x00, 0x00, 0x00, 0x3b, 0x00, 0x04, 0x00, 0x08, 0x00, 0x00, 0x00,
        0x09, 0x00, 0x00, 0x00, 0x03, 0x00, 0x00, 0x00, 0x1e, 0x00, 0x03, 0x00,
        0x0a, 0x00, 0x00, 0x00, 0x07, 0x00, 0x00, 0x00, 0x20, 0x00, 0x04, 0x00,
        0x0b, 0x00, 0x00, 0x00, 0x09, 0x00, 0x00, 0x00, 0x0a, 0x00, 0x00, 0x00,
        0x3b, 0x00, 0x04, 0x00, 0x0b, 0x00, 0x00, 0x00, 0x0c, 0x00, 0x00, 0x00,
        0x09, 0x00, 0x00, 0x00, 0x15, 0x00, 0x04, 0x00, 0x0d, 0x00, 0x00, 0x00,
        0x20, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x2b, 0x00, 0x04, 0x00,
        0x0d, 0x00, 0x00, 0x00, 0x0e, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
        0x20, 0x00, 0x04, 0x00, 0x0f, 0x00, 0x00, 0x00, 0x09, 0x00, 0x00, 0x00,
        0x07, 0x00, 0x00, 0x00, 0x36, 0x00, 0x05, 0x00, 0x02, 0x00, 0x00, 0x00,
        0x04, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x03, 0x00, 0x00, 0x00,
        0xf8, 0x00, 0x02, 0x00, 0x05, 0x00, 0x00, 0x00, 0x41, 0x00, 0x05, 0x00,
        0x0f, 0x00, 0x00, 0x00, 0x10, 0x00, 0x00, 0x00, 0x0c, 0x00, 0x00, 0x00,
        0x0e, 0x00, 0x00, 0x00, 0x3d, 0x00, 0x04, 0x00, 0x07, 0x00, 0x00, 0x00,
        0x11, 0x00, 0x00, 0x00, 0x10, 0x00, 0x00, 0x00, 0x3e, 0x00, 0x03, 0x00,
        0x09, 0x00, 0x00, 0x00, 0x11, 0x00, 0x00, 0x00, 0xfd, 0x00, 0x01, 0x00,
        0x38, 0x00, 0x01, 0x00
};
// clang-format on

// ---------------------------------------------------------------------------

class ExampleVulkanCustomLayer : public mbgl::style::CustomLayerHost {
public:
    ~ExampleVulkanCustomLayer() = default;

    void initialize(const mbgl::style::CustomLayerInitParameters &baseParams) override {
        const auto &params = static_cast<const mbgl::style::vulkan::CustomLayerInitParameters &>(baseParams);

        __android_log_print(ANDROID_LOG_INFO,
                            VK_LOG_TAG,
                            "initialize: device=%p, physicalDevice=%p",
                            static_cast<VkDevice>(params.device),
                            static_cast<VkPhysicalDevice>(params.physicalDevice));

        vulkanDevice = params.device;
        vulkanDispatcher = &params.dispatcher;

        createShaders(vulkanDevice, *vulkanDispatcher);
        createPipelineLayout(vulkanDevice, *vulkanDispatcher);
    }

    void render(const mbgl::style::CustomLayerRenderParameters &baseParams) override {
        const auto &params = static_cast<const mbgl::style::vulkan::CustomLayerRenderParameters &>(baseParams);

        const auto &dispatcher = params.dispatcher;
        const auto &cmd = params.commandBuffer;

        if (!pipeline) {
            createPipeline(params.device, dispatcher, params.renderPass);
        }

        if (!pipeline) return;

        cmd.pushConstants(
            *pipelineLayout, vk::ShaderStageFlagBits::eFragment, 0, sizeof(color), color.data(), dispatcher);

        cmd.bindPipeline(vk::PipelineBindPoint::eGraphics, *pipeline, dispatcher);

        cmd.draw(3, 1, 0, 0, dispatcher);

        __android_log_write(ANDROID_LOG_VERBOSE, VK_LOG_TAG, "render");
    }

    void contextLost() override {
        __android_log_write(ANDROID_LOG_INFO, VK_LOG_TAG, "contextLost");
        pipeline.reset();
        pipelineLayout.reset();
        vertexModule.reset();
        fragmentModule.reset();
    }

    void deinitialize() override {
        __android_log_write(ANDROID_LOG_INFO, VK_LOG_TAG, "deinitialize");
        pipeline.reset();
        pipelineLayout.reset();
        vertexModule.reset();
        fragmentModule.reset();
    }

    static std::array<float, 4> color;

private:
    void createShaders(const vk::Device &device, const vk::detail::DispatchLoaderDynamic &dispatcher) {
        try {
            vertexModule = device.createShaderModuleUnique(
                vk::ShaderModuleCreateInfo{.codeSize = sizeof(kVertSpirv),
                                           .pCode = reinterpret_cast<const std::uint32_t *>(kVertSpirv)},
                nullptr,
                dispatcher);
        } catch (const std::exception &) {
            __android_log_write(ANDROID_LOG_ERROR, VK_LOG_TAG, "Failed to create vertex shader module.");
            return;
        }

        try {
            fragmentModule = device.createShaderModuleUnique(
                vk::ShaderModuleCreateInfo{.codeSize = sizeof(kFragSpirv),
                                           .pCode = reinterpret_cast<const std::uint32_t *>(kFragSpirv)},
                nullptr,
                dispatcher);
        } catch (const std::exception &) {
            __android_log_write(ANDROID_LOG_ERROR, VK_LOG_TAG, "Failed to create fragment shader module.");
            return;
        }

        __android_log_write(ANDROID_LOG_INFO, VK_LOG_TAG, "Shader modules created");
    }

    void createPipelineLayout(const vk::Device &device, const vk::detail::DispatchLoaderDynamic &dispatcher) {
        const vk::PushConstantRange pushRange{
            .stageFlags = vk::ShaderStageFlagBits::eFragment,
            .offset = 0,
            .size = sizeof(float) * 4,
        };

        try {
            pipelineLayout = device.createPipelineLayoutUnique(
                vk::PipelineLayoutCreateInfo{
                    .pushConstantRangeCount = 1,
                    .pPushConstantRanges = &pushRange,
                },
                nullptr,
                dispatcher);
        } catch (const std::exception &) {
            __android_log_write(ANDROID_LOG_ERROR, VK_LOG_TAG, "Failed to create pipeline layout.");
            return;
        }

        __android_log_write(ANDROID_LOG_INFO, VK_LOG_TAG, "Pipeline layout created");
    }

    void createPipeline(const vk::Device &device,
                        const vk::detail::DispatchLoaderDynamic &dispatcher,
                        vk::RenderPass renderPass) {
        if (!vertexModule || !fragmentModule || !pipelineLayout) return;

        const std::array<vk::PipelineShaderStageCreateInfo, 2> stages = {{
            {
                .stage = vk::ShaderStageFlagBits::eVertex,
                .module = *vertexModule,
                .pName = "main",
            },
            {
                .stage = vk::ShaderStageFlagBits::eFragment,
                .module = *fragmentModule,
                .pName = "main",
            },
        }};

        const vk::PipelineVertexInputStateCreateInfo vertexInput{};

        const vk::PipelineInputAssemblyStateCreateInfo inputAssembly{
            .topology = vk::PrimitiveTopology::eTriangleList,
        };

        const std::array<vk::DynamicState, 2> dynamicStates = {
            vk::DynamicState::eViewport,
            vk::DynamicState::eScissor,
        };
        const vk::PipelineDynamicStateCreateInfo dynamicState{
            .dynamicStateCount = static_cast<uint32_t>(dynamicStates.size()),
            .pDynamicStates = dynamicStates.data(),
        };

        const vk::PipelineViewportStateCreateInfo viewportState{
            .viewportCount = 1,
            .scissorCount = 1,
        };

        const vk::PipelineRasterizationStateCreateInfo rasterizer{
            .polygonMode = vk::PolygonMode::eFill,
            .cullMode = vk::CullModeFlagBits::eNone,
            .frontFace = vk::FrontFace::eClockwise,
            .lineWidth = 1.0f,
        };

        const vk::PipelineMultisampleStateCreateInfo multisampling{
            .rasterizationSamples = vk::SampleCountFlagBits::e1,
        };

        const vk::PipelineColorBlendAttachmentState blendAttachment{
            .blendEnable = VK_TRUE,
            .srcColorBlendFactor = vk::BlendFactor::eSrcAlpha,
            .dstColorBlendFactor = vk::BlendFactor::eOneMinusSrcAlpha,
            .colorBlendOp = vk::BlendOp::eAdd,
            .srcAlphaBlendFactor = vk::BlendFactor::eOne,
            .dstAlphaBlendFactor = vk::BlendFactor::eZero,
            .alphaBlendOp = vk::BlendOp::eAdd,
            .colorWriteMask = vk::ColorComponentFlagBits::eR | vk::ColorComponentFlagBits::eG |
                              vk::ColorComponentFlagBits::eB | vk::ColorComponentFlagBits::eA,
        };

        const vk::PipelineColorBlendStateCreateInfo colorBlend{
            .attachmentCount = 1,
            .pAttachments = &blendAttachment,
        };

        const vk::PipelineDepthStencilStateCreateInfo depthStencil{
            .depthTestEnable = VK_FALSE,
            .depthWriteEnable = VK_FALSE,
        };

        const vk::GraphicsPipelineCreateInfo pipelineInfo{
            .stageCount = static_cast<uint32_t>(stages.size()),
            .pStages = stages.data(),
            .pVertexInputState = &vertexInput,
            .pInputAssemblyState = &inputAssembly,
            .pViewportState = &viewportState,
            .pRasterizationState = &rasterizer,
            .pMultisampleState = &multisampling,
            .pDepthStencilState = &depthStencil,
            .pColorBlendState = &colorBlend,
            .pDynamicState = &dynamicState,
            .layout = *pipelineLayout,
            .renderPass = renderPass,
            .subpass = 0,
        };

        auto result = device.createGraphicsPipelineUnique(nullptr, pipelineInfo, nullptr, dispatcher);
        if (result.result != vk::Result::eSuccess) {
            __android_log_write(ANDROID_LOG_ERROR, VK_LOG_TAG, "Failed to create graphics pipeline");
            return;
        }
        pipeline = std::move(result.value);
        __android_log_write(ANDROID_LOG_INFO, VK_LOG_TAG, "Pipeline created");
    }

    vk::Device vulkanDevice;
    const vk::detail::DispatchLoaderDynamic *vulkanDispatcher = nullptr;

    vk::UniqueHandle<vk::ShaderModule, vk::detail::DispatchLoaderDynamic> vertexModule;
    vk::UniqueHandle<vk::ShaderModule, vk::detail::DispatchLoaderDynamic> fragmentModule;

    vk::UniqueHandle<vk::Pipeline, vk::detail::DispatchLoaderDynamic> pipeline;
    vk::UniqueHandle<vk::PipelineLayout, vk::detail::DispatchLoaderDynamic> pipelineLayout;
};

std::array<float, 4> ExampleVulkanCustomLayer::color = {0.0f, 1.0f, 0.0f, 0.5f};

// ---------------------------------------------------------------------------
// JNI entry points
// ---------------------------------------------------------------------------

jlong JNICALL nativeCreateContext(JNIEnv *, jobject) {
    __android_log_write(ANDROID_LOG_INFO, VK_LOG_TAG, "nativeCreateContext");
    return reinterpret_cast<jlong>(new ExampleVulkanCustomLayer());
}

void JNICALL nativeSetColor(JNIEnv *, jobject, jfloat red, jfloat green, jfloat blue, jfloat alpha) {
    __android_log_print(ANDROID_LOG_INFO,
                        VK_LOG_TAG,
                        "ExampleVulkanCustomLayer.nativeSetColor: %.2f, %.2f, %.2f, %.2f",
                        red,
                        green,
                        blue,
                        alpha);
    ExampleVulkanCustomLayer::color = {red, green, blue, alpha};
}

extern "C" JNIEXPORT jint JNICALL JNI_OnLoad(JavaVM *vm, void *) {
    __android_log_write(ANDROID_LOG_INFO, VK_LOG_TAG, "JNI_OnLoad");

    JNIEnv *env = nullptr;
    vm->GetEnv(reinterpret_cast<void **>(&env), JNI_VERSION_1_6);

    jclass cls = env->FindClass("org/maplibre/android/testapp/model/customlayer/ExampleVulkanCustomLayer");

    const JNINativeMethod methods[] = {
        {"createContext", "()J", reinterpret_cast<void *>(&nativeCreateContext)},
        {"setColor", "(FFFF)V", reinterpret_cast<void *>(&nativeSetColor)},
    };

    if (env->RegisterNatives(cls, methods, 2) < 0) {
        env->ExceptionDescribe();
        return JNI_ERR;
    }

    return JNI_VERSION_1_6;
}

extern "C" JNIEXPORT void JNICALL JNI_OnUnload(JavaVM *, void *) {}