#include "OpenGLGfxInterface.h"
#include "Log.h"
#include "AssetManager.h"
#include "Texture.h"
#include "Profiling.h"

#ifdef _WIN32
    #include "glad.h"
#elif defined(__ANDROID__)
    #include <GLES3/gl3.h>
#else
    #pragma error "Don't compile this!"
#endif
#include <unordered_map>
#include <vector>

#include <assert.h>


static const char* g_vert_src = "#version 300 es                                                                   \n"
                                "precision mediump float;                                                          \n"
                                "uniform vec2 res; // resolution                                                   \n"
                                "layout (location = 0) in vec2 v_p0;   // top left corner on screen                \n"
                                "layout (location = 1) in vec2 v_p1;   // top right corner on screen               \n"
                                "layout (location = 2) in vec2 v_src_p0;                                           \n"
                                "layout (location = 3) in vec2 v_src_p1;                                           \n"
                                "layout (location = 4) in vec4 v_color;    // color                                \n"
                                "layout (location = 5) in int v_expand_r; //                                       \n"
                                "out vec4 f_color;                                                                 \n"
                                "out vec2 f_uv;                                                                    \n"
                                "flat out int f_expand_r;                                                          \n"
                                "const vec2 vertices[6] = vec2[6](                                                 \n"
                                "        vec2(-1,-1),                                                              \n"
                                "        vec2(+1, -1),                                                             \n"
                                "        vec2(-1, +1),                                                             \n"
                                "        vec2(+1, -1),                                                             \n"
                                "        vec2(+1, +1),                                                             \n"
                                "        vec2(-1, +1)                                                              \n"
                                ");                                                                                \n"
                                "void main()                                                                       \n"
                                "{                                                                                 \n"
                                "    vec2 v = vertices[gl_VertexID];                                               \n"
                                "    // destination on screen                                                      \n"
                                "    //vec2 dst_half_size = (v_p1 - v_p0) / 2.0;                                   \n"
                                "    //vec2 dst_center = (v_p1 + v_p0) / 2.0;                                      \n"
                                "    vec2 dst_pos = v * ((v_p1 - v_p0) / 2.0) + ((v_p1 + v_p0) / 2.0);             \n"
                                "    dst_pos.y = res.y - dst_pos.y;                                                \n"
                                "    vec2 src = v * ((v_src_p1 - v_src_p0) / 2.0) + ((v_src_p1 + v_src_p0) / 2.0); \n"
                                "    gl_Position = vec4(2.0 * dst_pos.x / res.x - 1.0,                             \n"
                                "                       2.0 * dst_pos.y / res.y - 1.0,                             \n"
                                "                       0.0,                                                       \n"
                                "                       1.0);                                                      \n"
                                "    f_color = v_color;                                                            \n"
                                "    f_uv = src;                                                                   \n"
                                "    f_expand_r = v_expand_r;                                                      \n"
                                "}                                                                                 \n";

static const char* g_frag_src = "#version 300 es                                      \n"
                                "precision mediump float;                             \n"
                                "in vec2 f_uv;                                        \n"
                                "in vec4 f_color;                                     \n"
                                "flat in int f_expand_r;                              \n"
                                "layout (location = 0) out vec4 frag_color;           \n"
                                "uniform sampler2D s_texture;                         \n"
                                "void main()                                          \n"
                                "{                                                    \n"
                                "    vec4 tex_color = texture(s_texture, f_uv);       \n"
                                "    if (f_expand_r == 1)                             \n"
                                "        tex_color = tex_color.rrrr;                  \n"
                                "    frag_color = tex_color * f_color;                \n"
                                "}                                                    \n";

struct OpenGLGfx
{
    GLuint m_shader;
    GLuint m_vao;
    GLuint m_vbo;

    KDEStringHandle m_dummy_texture;

    std::vector<KDERect> m_rects;
    std::vector<KDEStringHandle> m_draw_textures;

    std::unordered_map<KDEStringHandle, Texture> m_textures;
}

OpenGLGfx* createOpenGLGfx();
{
    ZoneScoped;
    OpenGLGFX* gfx = new OpenGLGFX;
    
    GLuint vertex = 0, fragment = 0;
    vertex = glCreateShader(GL_VERTEX_SHADER);
    glShaderSource(vertex, 1, &g_vert_src, nullptr);
    glCompileShader(vertex);
    int status = 0;
    glGetShaderiv(vertex, GL_COMPILE_STATUS, &status);
    if (status != GL_TRUE) {
        char info[512];
        glGetShaderInfoLog(vertex, 512, nullptr, info);
        ALOGE("Failed to compile ui vertex shader: %s", info);
    }

    fragment = glCreateShader(GL_FRAGMENT_SHADER);
    glShaderSource(fragment, 1, &g_frag_src, nullptr);
    glCompileShader(fragment);
    glGetShaderiv(fragment, GL_COMPILE_STATUS, &status);
    if (status != GL_TRUE) {
        char info[512];
        glGetShaderInfoLog(fragment, 512, nullptr, info);
        ALOGE("Failed to compile ui fragment shader: %s", info);
    }

    gfx->m_shader = glCreateProgram();
    glAttachShader(gfx->m_shader, vertex);
    glAttachShader(gfx->m_shader, fragment);
    glLinkProgram(gfx->m_shader);
    glGetProgramiv(gfx->m_shader, GL_LINK_STATUS, &status);
    if (status != GL_TRUE) {
        char info[512];
        glGetProgramInfoLog(gfx->m_shader, 512, nullptr, info);
        ALOGE("Failed to link ui shader: %s", info);
    }

    glDeleteShader(vertex);
    glDeleteShader(fragment);

    glGenBuffers(1, &gfx->m_vbo);
    glGenVertexArrays(1, &gfx->m_vao);
    glBindVertexArray(gfx->m_vao);
    glBindBuffer(GL_ARRAY_BUFFER, gfx->m_vbo);

    glVertexAttribPointer(0, 2, GL_FLOAT, GL_FALSE, sizeof(Rect), (GLvoid*)0);
    glVertexAttribPointer(1, 2, GL_FLOAT, GL_FALSE, sizeof(Rect), (GLvoid*)offsetof(Rect, dst_p1_x));
    glVertexAttribPointer(2, 2, GL_FLOAT, GL_FALSE, sizeof(Rect), (GLvoid*)offsetof(Rect, src_p0_x));
    glVertexAttribPointer(3, 2, GL_FLOAT, GL_FALSE, sizeof(Rect), (GLvoid*)offsetof(Rect, src_p1_x));
    glVertexAttribPointer(4, 4, GL_FLOAT, GL_FALSE, sizeof(Rect), (GLvoid*)offsetof(Rect, r));
    glVertexAttribIPointer(5, 1, GL_INT, sizeof(Rect), (GLvoid*)offsetof(Rect, expand_r));
    glEnableVertexAttribArray(0);
    glEnableVertexAttribArray(1);
    glEnableVertexAttribArray(2);
    glEnableVertexAttribArray(3);
    glEnableVertexAttribArray(4);
    glEnableVertexAttribArray(5);
    glVertexAttribDivisor(0, 1);
    glVertexAttribDivisor(1, 1);
    glVertexAttribDivisor(2, 1);
    glVertexAttribDivisor(3, 1);
    glVertexAttribDivisor(4, 1);
    glVertexAttribDivisor(5, 1);
    glBindVertexArray(0);

    // Dummy texture (1x1px white)
    unsigned char tex[] = {255, 255, 255, 255};
    Texture dummy_texture(1, 1, tex);
    gfx->m_textures.insert(std::make_pair(0, dummy_texture));
    m_dummy_texture = 0;
}

void destroyOpenGLGfx(OpenGLGFX* gfx)
{
    ZoneScoped;
    for (auto texture : gfx->m_textures)
        texture.second.destroy();
    glDeleteProgram(gfx->m_shader);
    glDeleteVertexArrays(1, &gfx->m_vao);
    glDeleteBuffers(1, &gfx->m_vbo);
    
    delete gfx;
}

void openGLGfxPushRect(OpenGLGfx* gfx, KDERect rect, KDEStringHandle texture)
{
    if (gfx->m_textures.find(texture) == gfx->m_textures.end()) {
        Texture tex;
        if (!AssetManager::ptr->loadTexture(StringRepository::global->getString(texture), &tex)) {
            ALOGE("Failed to load texture");
            return;
        }
        gfx->m_textures.insert(std::make_pair(texture, tex));
    }
    
    gfx->m_rects.push_back(rect);
    gfx->m_draw_textures.push_back(texture);
}


void openGLGfxPushFontRect(OpenGLGfx* gfx, KDERect rect,
                           float char_height,
                           KDEStringHandle font,
                           KDEStringHandle id)
{
    if (gfx->m_textures.find(id) == gfx->m_textures.end()) {
        Texture tex;
        if (!AssetManager::ptr->loadFontBitmap(StringRepository::global->getString(font), char_height, &tex)) {
            ALOGE("Failed to load font");
            return;
        }
        gfx->m_textures.insert(std::make_pair(id, tex));
    }
    
    gfx->m_rects.push_back(rect);
    gfx->m_draw_textures.push_back(texture);
}


void openGLGetTextureSize(OpenGLGfx* gfx, KDEStringHandle texture, unsigned int* w, unsigned int* h)
{
    ZoneScoped;
    if (gfx->m_textures.find(texture) == gfx->m_textures.end()) {
        Texture tex;
        if (!AssetManager::ptr->loadTexture(StringRepository::global->getString(texture), &tex)) {
            ALOGE("Failed to load texture");
            return;
        }
        gfx->m_textures.insert(std::make_pair(texture, tex));
    }
    Texture tex = gfx->m_textures[texture];
    tex.getTextureSize(w, h);
}

void openGLRenderFrame(OpenGLGfx* gfx, float width, float height)
{
    ZoneScoped;
    assert(gfx->m_rects.size() == gfx->m_draw_textures.size());

    glViewport(0, 0, static_cast<int>(width), static_cast<int>(height));
    // glClearColor(0.8f, 0.3f, 0.3f, 1.f);
    glClearColor(0.0f, 0.0f, 0.0f, 1.f);
    glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);

    const auto rect_count = static_cast<GLsizei>(gfx->m_rects.size());
    if (rect_count == 0)
        return;

    glDisable(GL_DEPTH_TEST);
    glDisable(GL_CULL_FACE);
    glEnable(GL_BLEND);
    glBlendEquation(GL_FUNC_ADD);
    glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA);

    glBindVertexArray(gfx->m_vao);
    glUseProgram(gfx->m_shader);
    glUniform2f(glGetUniformLocation(gfx->m_shader, "res"), width, height);
    glUniform1i(glGetUniformLocation(gfx->m_shader, "s_texture"), 0);
    glBindBuffer(GL_ARRAY_BUFFER, gfx->m_vbo);

    size_t start = 0;
    StringHandle current_texture = gfx->m_draw_textures[0];
    while (start < rect_count) {
        glActiveTexture(GL_TEXTURE0);
        gfx->m_textures[current_texture].bind();
        for (size_t i = start; i < rect_count; ++i) {
            StringHandle texture = gfx->m_draw_textures[i];
            if ((texture != current_texture)) {
                // Draw everything in [start, i - 1]
                const size_t count = i - start;
                glBufferData(GL_ARRAY_BUFFER,
                             static_cast<GLsizeiptr>(sizeof(Rect) * count),
                             &gfx->m_rects[start],
                             GL_STREAM_DRAW);
                glDrawArraysInstanced(GL_TRIANGLES, 0, 6, static_cast<GLsizei>(count));

                start = i;
                current_texture = texture;
                break;
            }

            if (i == rect_count - 1) {
                // just draw it
                const size_t count = rect_count - start;
                glBufferData(GL_ARRAY_BUFFER,
                             static_cast<GLsizeiptr>(sizeof(Rect) * count),
                             &gfx->m_rects[start],
                             GL_STREAM_DRAW);
                glDrawArraysInstanced(GL_TRIANGLES, 0, 6, static_cast<GLsizei>(count));

                start = rect_count;
                current_texture = texture;
            }
        }
    }
    gfx->m_rects.clear();
    gfx->m_draw_textures.clear();
}
