#include "stb_truetype.h"
#include "stb_image.h"

#include "vui_internal.h"

namespace vui {

Renderer::Renderer() : impl(new Impl) {}
Renderer::~Renderer() { delete impl; }

void Renderer::init(VkDevice device, VkPhysicalDevice physDevice,
                    VkQueue graphicsQueue, uint32_t queueFamily,
                    VkFormat swapchainFormat, uint32_t imageCount) {
    impl->device = device;
    impl->physDevice = physDevice;
    impl->graphicsQueue = graphicsQueue;
    impl->queueFamily = queueFamily;
    impl->swapchainFormat = swapchainFormat;
    impl->imageCount = imageCount;

    impl->initRenderPass();
    impl->initSampler();
    impl->initDescriptorPool();
    impl->initPipeline();
    impl->initVertexBuffer(INITIAL_VERTEX_COUNT);
    impl->createWhiteTexture();

    impl->currentTransform = Mat3::identity();
    memset(&impl->input, 0, sizeof(impl->input));
}

void Renderer::shutdown() {
    vkDeviceWaitIdle(impl->device);

    for (auto &tex : impl->textures) {
        if (!tex.active) continue;
        vkDestroyImageView(impl->device, tex.view, nullptr);
        vkDestroyImage(impl->device, tex.image, nullptr);
        vkFreeMemory(impl->device, tex.memory, nullptr);
    }
    impl->textures.clear();
    impl->fonts.clear();

    if (impl->vertexBuffer) {
        vkDestroyBuffer(impl->device, impl->vertexBuffer, nullptr);
        vkFreeMemory(impl->device, impl->vertexMemory, nullptr);
    }

    for (auto fb : impl->overlayFramebuffers)
        vkDestroyFramebuffer(impl->device, fb, nullptr);
    impl->overlayFramebuffers.clear();

    if (impl->pipeline) vkDestroyPipeline(impl->device, impl->pipeline, nullptr);
    if (impl->pipelineLayout) vkDestroyPipelineLayout(impl->device, impl->pipelineLayout, nullptr);
    if (impl->descriptorPool) vkDestroyDescriptorPool(impl->device, impl->descriptorPool, nullptr);
    if (impl->descriptorSetLayout) vkDestroyDescriptorSetLayout(impl->device, impl->descriptorSetLayout, nullptr);
    if (impl->sampler) vkDestroySampler(impl->device, impl->sampler, nullptr);
    if (impl->renderPass) vkDestroyRenderPass(impl->device, impl->renderPass, nullptr);
}

void Renderer::setOverlayFramebuffers(VkImageView *views, uint32_t count, uint32_t width, uint32_t height) {
    for (auto fb : impl->overlayFramebuffers)
        vkDestroyFramebuffer(impl->device, fb, nullptr);
    impl->overlayFramebuffers.clear();

    impl->overlayWidth = width;
    impl->overlayHeight = height;
    impl->overlayFramebuffers.resize(count);

    for (uint32_t i = 0; i < count; i++) {
        VkFramebufferCreateInfo fci{VK_STRUCTURE_TYPE_FRAMEBUFFER_CREATE_INFO};
        fci.renderPass = impl->renderPass;
        fci.attachmentCount = 1;
        fci.pAttachments = &views[i];
        fci.width = width;
        fci.height = height;
        fci.layers = 1;
        vkCreateFramebuffer(impl->device, &fci, nullptr, &impl->overlayFramebuffers[i]);
    }
}

VkRenderPass Renderer::getRenderPass() const {
    return impl->renderPass;
}

void Renderer::beginFrame(int windowW, int windowH, int framebufferW, int framebufferH) {
    impl->winW = windowW;
    impl->winH = windowH;
    impl->fbW = (framebufferW > 0) ? framebufferW : windowW;
    impl->fbH = (framebufferH > 0) ? framebufferH : windowH;
    impl->vertices.clear();
    impl->drawCmds.clear();
    impl->transformStack.clear();
    impl->scissorStack.clear();
    impl->currentTransform = Mat3::identity();
    impl->frameActive = true;

    memcpy(impl->input.prevMouseButtons, impl->input.mouseButtons, sizeof(impl->input.mouseButtons));
    memcpy(impl->input.prevKeys, impl->input.keys, sizeof(impl->input.keys));
    memset(impl->input.keyClicked, 0, sizeof(impl->input.keyClicked));
    memset(impl->input.mouseClicked, 0, sizeof(impl->input.mouseClicked));
    impl->input.scrollX = 0;
    impl->input.scrollY = 0;
    impl->input.charCount = 0;
}

void Renderer::endFrame(VkCommandBuffer cmd, uint32_t imageIndex) {
    impl->frameActive = false;

    if (impl->vertices.empty()) return;

    if (impl->vertices.size() > impl->vertexBufferCapacity) {
        uint32_t newCap = (uint32_t)impl->vertices.size() * 2;
        impl->initVertexBuffer(newCap);
    }

    void *mapped;
    vkMapMemory(impl->device, impl->vertexMemory, 0,
                sizeof(Vertex) * impl->vertices.size(), 0, &mapped);
    memcpy(mapped, impl->vertices.data(), sizeof(Vertex) * impl->vertices.size());
    vkUnmapMemory(impl->device, impl->vertexMemory);

    bool ownPass = !impl->overlayFramebuffers.empty() &&
                   imageIndex < (uint32_t)impl->overlayFramebuffers.size();

    if (ownPass) {
        VkRenderPassBeginInfo rpbi{VK_STRUCTURE_TYPE_RENDER_PASS_BEGIN_INFO};
        rpbi.renderPass = impl->renderPass;
        rpbi.framebuffer = impl->overlayFramebuffers[imageIndex];
        rpbi.renderArea.extent = {impl->overlayWidth, impl->overlayHeight};
        vkCmdBeginRenderPass(cmd, &rpbi, VK_SUBPASS_CONTENTS_INLINE);
    }

    float L = 0, R = (float)impl->winW, T = 0, B = (float)impl->winH;
    float projection[16] = {
        2.0f/(R-L),    0.0f,           0.0f, 0.0f,
        0.0f,          2.0f/(B-T),     0.0f, 0.0f,
        0.0f,          0.0f,          -1.0f, 0.0f,
        -(R+L)/(R-L), -(B+T)/(B-T),   0.0f, 1.0f,
    };

    vkCmdBindPipeline(cmd, VK_PIPELINE_BIND_POINT_GRAPHICS, impl->pipeline);
    vkCmdPushConstants(cmd, impl->pipelineLayout, VK_SHADER_STAGE_VERTEX_BIT,
                       0, sizeof(projection), projection);

    VkViewport viewport{0, 0, (float)impl->fbW, (float)impl->fbH, 0.0f, 1.0f};
    vkCmdSetViewport(cmd, 0, 1, &viewport);

    VkRect2D fullScissor{{0,0}, {(uint32_t)impl->fbW, (uint32_t)impl->fbH}};
    vkCmdSetScissor(cmd, 0, 1, &fullScissor);

    VkDeviceSize offset = 0;
    vkCmdBindVertexBuffers(cmd, 0, 1, &impl->vertexBuffer, &offset);

    float scaleX = (float)impl->fbW / (float)impl->winW;
    float scaleY = (float)impl->fbH / (float)impl->winH;

    int lastTexId = -1;
    for (auto &dc : impl->drawCmds) {
        if (dc.textureId != lastTexId) {
            if (dc.textureId >= 0 && dc.textureId < (int)impl->textures.size()) {
                vkCmdBindDescriptorSets(cmd, VK_PIPELINE_BIND_POINT_GRAPHICS,
                                        impl->pipelineLayout, 0, 1,
                                        &impl->textures[dc.textureId].descriptorSet,
                                        0, nullptr);
            }
            lastTexId = dc.textureId;
        }

        if (dc.hasScissor) {
            VkRect2D sc;
            sc.offset.x = std::max(0, (int)(dc.scissorX * scaleX));
            sc.offset.y = std::max(0, (int)(dc.scissorY * scaleY));
            sc.extent.width = (uint32_t)std::max(0.0f, dc.scissorW * scaleX);
            sc.extent.height = (uint32_t)std::max(0.0f, dc.scissorH * scaleY);
            vkCmdSetScissor(cmd, 0, 1, &sc);
        } else {
            vkCmdSetScissor(cmd, 0, 1, &fullScissor);
        }

        vkCmdDraw(cmd, dc.vertexCount, 1, dc.vertexOffset, 0);
    }

    if (ownPass) {
        vkCmdEndRenderPass(cmd);
    }
}

void Renderer::drawRect(float x, float y, float w, float h, Color color) {
    impl->addQuad(WHITE_TEXTURE_ID, x, y, x+w, y+h, 0, 0, 1, 1, color);
}

void Renderer::drawRectOutline(float x, float y, float w, float h, float t, Color color) {
    drawRect(x, y, w, t, color);
    drawRect(x, y+h-t, w, t, color);
    drawRect(x, y+t, t, h-2*t, color);
    drawRect(x+w-t, y+t, t, h-2*t, color);
}

void Renderer::drawImage(int texId, float x, float y, float w, float h, Color tint) {
    impl->addQuad(texId, x, y, x+w, y+h, 0, 0, 1, 1, tint);
}

void Renderer::drawImageRegion(int texId, float x, float y, float w, float h,
                               float u0, float v0, float u1, float v1, Color tint) {
    impl->addQuad(texId, x, y, x+w, y+h, u0, v0, u1, v1, tint);
}

void Renderer::drawImage9Slice(int texId, float x, float y, float w, float h,
                               float bL, float bR, float bT, float bB,
                               int texW, int texH, Color tint) {
    float uL = bL / texW, uR = 1.0f - bR / texW;
    float vT = bT / texH, vB = 1.0f - bB / texH;
    float midW = w - bL - bR, midH = h - bT - bB;

    drawImageRegion(texId, x,          y,          bL,   bT,   0,  0,  uL, vT, tint);
    drawImageRegion(texId, x+bL,       y,          midW, bT,   uL, 0,  uR, vT, tint);
    drawImageRegion(texId, x+bL+midW,  y,          bR,   bT,   uR, 0,  1,  vT, tint);

    drawImageRegion(texId, x,          y+bT,       bL,   midH, 0,  vT, uL, vB, tint);
    drawImageRegion(texId, x+bL,       y+bT,       midW, midH, uL, vT, uR, vB, tint);
    drawImageRegion(texId, x+bL+midW,  y+bT,       bR,   midH, uR, vT, 1,  vB, tint);

    drawImageRegion(texId, x,          y+bT+midH,  bL,   bB,   0,  vB, uL, 1,  tint);
    drawImageRegion(texId, x+bL,       y+bT+midH,  midW, bB,   uL, vB, uR, 1,  tint);
    drawImageRegion(texId, x+bL+midW,  y+bT+midH,  bR,   bB,   uR, vB, 1,  1,  tint);
}

void Renderer::drawText(const char *text, float x, float y, Color color, int fontId) {
    if (fontId < 0 || fontId >= (int)impl->fonts.size() || !impl->fonts[fontId].active) return;
    auto &font = impl->fonts[fontId];
    float cx = x;
    float cy = y + font.ascent;

    for (const char *p = text; *p; p++) {
        int ch = (unsigned char)*p;
        if (ch >= 128) continue;
        auto &g = font.glyphs[ch];
        if (g.x1 > g.x0) {
            float gx = cx + g.xOff;
            float gy = cy + g.yOff;
            impl->addQuad(font.textureId,
                          gx, gy, gx + (g.x1 - g.x0), gy + (g.y1 - g.y0),
                          g.u0, g.v0, g.u1, g.v1, color);
        }
        cx += g.xAdvance;
    }
}

void Renderer::drawTextCentered(const char *text, float x, float y, float w, float h,
                                Color color, int fontId) {
    float tw = measureText(text, fontId);
    float th = fontHeight(fontId);
    float tx = x + (w - tw) * 0.5f;
    float ty = y + (h - th) * 0.5f;
    drawText(text, tx, ty, color, fontId);
}

void Renderer::drawTextWrapped(const char *text, float x, float y, float maxWidth,
                               Color color, int fontId) {
    if (fontId < 0 || fontId >= (int)impl->fonts.size() || !impl->fonts[fontId].active) return;
    auto &font = impl->fonts[fontId];

    float cx = x;
    float cy = y;
    float lineH = font.ascent - font.descent + font.lineGap;
    const char *wordStart = text;

    while (*wordStart) {
        while (*wordStart == ' ') {
            int ch = (unsigned char)*wordStart;
            if (ch < 128) cx += font.glyphs[ch].xAdvance;
            wordStart++;
        }
        if (!*wordStart) break;

        const char *wordEnd = wordStart;
        float wordW = 0;
        while (*wordEnd && *wordEnd != ' ') {
            int ch = (unsigned char)*wordEnd;
            if (ch < 128) wordW += font.glyphs[ch].xAdvance;
            wordEnd++;
        }

        if (cx + wordW > x + maxWidth && cx > x) {
            cx = x;
            cy += lineH;
        }

        for (const char *p = wordStart; p < wordEnd; p++) {
            int ch = (unsigned char)*p;
            if (ch < 128) {
                auto &g = font.glyphs[ch];
                if (g.x1 > g.x0) {
                    float gx = cx + g.xOff;
                    float gy = cy + font.ascent + g.yOff;
                    impl->addQuad(font.textureId,
                                  gx, gy, gx + (g.x1 - g.x0), gy + (g.y1 - g.y0),
                                  g.u0, g.v0, g.u1, g.v1, color);
                }
                cx += g.xAdvance;
            }
        }
        wordStart = wordEnd;
    }
}

float Renderer::measureText(const char *text, int fontId) {
    if (fontId < 0 || fontId >= (int)impl->fonts.size() || !impl->fonts[fontId].active) return 0;
    auto &font = impl->fonts[fontId];
    float w = 0;
    for (const char *p = text; *p; p++) {
        int ch = (unsigned char)*p;
        if (ch < 128) w += font.glyphs[ch].xAdvance;
    }
    return w;
}

float Renderer::fontHeight(int fontId) {
    if (fontId < 0 || fontId >= (int)impl->fonts.size() || !impl->fonts[fontId].active) return 0;
    auto &f = impl->fonts[fontId];
    return f.ascent - f.descent;
}

int Renderer::createTexture(int width, int height, const void *rgba) {
    return impl->allocTexture(width, height, rgba);
}

void Renderer::updateTexture(int id, int x, int y, int w, int h, const void *rgba) {
    if (id < 0 || id >= (int)impl->textures.size() || !impl->textures[id].active) return;
    impl->uploadTextureData(impl->textures[id], x, y, w, h, rgba);
}

void Renderer::destroyTexture(int id) {
    if (id < 0 || id >= (int)impl->textures.size() || !impl->textures[id].active) return;
    auto &tex = impl->textures[id];
    vkDeviceWaitIdle(impl->device);
    vkDestroyImageView(impl->device, tex.view, nullptr);
    vkDestroyImage(impl->device, tex.image, nullptr);
    vkFreeMemory(impl->device, tex.memory, nullptr);
    tex.active = false;
}

int Renderer::loadTexture(const char *path) {
    int w, h, channels;
    unsigned char *data = stbi_load(path, &w, &h, &channels, 4);
    if (!data) {
        fprintf(stderr, "vui: failed to load texture '%s'\n", path);
        return -1;
    }
    int id = createTexture(w, h, data);
    stbi_image_free(data);
    return id;
}

int Renderer::loadFont(const char *path, float size) {
    FILE *f = fopen(path, "rb");
    if (!f) {
        fprintf(stderr, "vui: failed to open font '%s'\n", path);
        return -1;
    }
    fseek(f, 0, SEEK_END);
    long len = ftell(f);
    fseek(f, 0, SEEK_SET);
    std::vector<unsigned char> buf(len);
    fread(buf.data(), 1, len, f);
    fclose(f);
    return loadFontFromMemory(buf.data(), (int)len, size);
}

int Renderer::loadFontFromMemory(const void *data, int dataSize, float size) {
    (void)dataSize;
    stbtt_fontinfo info;
    if (!stbtt_InitFont(&info, (const unsigned char *)data, 0)) {
        fprintf(stderr, "vui: failed to init font\n");
        return -1;
    }

    float fontScale = stbtt_ScaleForPixelHeight(&info, size);
    int ascent, descent, lineGap;
    stbtt_GetFontVMetrics(&info, &ascent, &descent, &lineGap);

    int atlasW = 512, atlasH = 512;
    std::vector<unsigned char> bitmap(atlasW * atlasH, 0);

    stbtt_bakedchar cdata[128];
    stbtt_BakeFontBitmap((const unsigned char *)data, 0, size,
                         bitmap.data(), atlasW, atlasH, 0, 128, cdata);

    std::vector<uint32_t> rgba(atlasW * atlasH);
    for (int i = 0; i < atlasW * atlasH; i++) {
        uint8_t a = bitmap[i];
        rgba[i] = (uint32_t)a << 24 | 0x00FFFFFF;
    }

    int texId = createTexture(atlasW, atlasH, rgba.data());

    FontEntry fe{};
    fe.textureId = texId;
    fe.size = size;
    fe.ascent = ascent * fontScale;
    fe.descent = descent * fontScale;
    fe.lineGap = lineGap * fontScale;
    fe.active = true;

    for (int i = 0; i < 128; i++) {
        auto &bc = cdata[i];
        fe.glyphs[i].x0 = (float)bc.x0;
        fe.glyphs[i].y0 = (float)bc.y0;
        fe.glyphs[i].x1 = (float)bc.x1;
        fe.glyphs[i].y1 = (float)bc.y1;
        fe.glyphs[i].u0 = (float)bc.x0 / atlasW;
        fe.glyphs[i].v0 = (float)bc.y0 / atlasH;
        fe.glyphs[i].u1 = (float)bc.x1 / atlasW;
        fe.glyphs[i].v1 = (float)bc.y1 / atlasH;
        fe.glyphs[i].xAdvance = bc.xadvance;
        fe.glyphs[i].xOff = bc.xoff;
        fe.glyphs[i].yOff = bc.yoff;
    }

    int id = (int)impl->fonts.size();
    impl->fonts.push_back(fe);
    return id;
}

void Renderer::pushScissor(float x, float y, float w, float h) {
    impl->scissorStack.push_back({x, y, w, h});
}

void Renderer::popScissor() {
    if (!impl->scissorStack.empty())
        impl->scissorStack.pop_back();
}

void Renderer::pushTransform() {
    impl->transformStack.push_back(impl->currentTransform);
}

void Renderer::popTransform() {
    if (!impl->transformStack.empty()) {
        impl->currentTransform = impl->transformStack.back();
        impl->transformStack.pop_back();
    }
}

void Renderer::translate(float x, float y) {
    Mat3 t;
    t.m[2] = x;
    t.m[5] = y;
    impl->currentTransform = impl->currentTransform * t;
}

void Renderer::rotate(float angleDeg) {
    float rad = angleDeg * 3.14159265f / 180.0f;
    float c = cosf(rad), s = sinf(rad);
    Mat3 r;
    r.m[0] = c; r.m[1] = -s;
    r.m[3] = s; r.m[4] = c;
    impl->currentTransform = impl->currentTransform * r;
}

void Renderer::scale(float sx, float sy) {
    Mat3 s;
    s.m[0] = sx;
    s.m[4] = sy;
    impl->currentTransform = impl->currentTransform * s;
}

void Renderer::feedMousePosition(float x, float y) {
    impl->input.mouseX = x;
    impl->input.mouseY = y;
}

void Renderer::feedMouseButton(int button, bool pressed) {
    if (button >= 0 && button < 5) {
        if (pressed && !impl->input.mouseButtons[button])
            impl->input.mouseClicked[button] = true;
        impl->input.mouseButtons[button] = pressed;
    }
}

void Renderer::feedMouseScroll(float dx, float dy) {
    impl->input.scrollX += dx;
    impl->input.scrollY += dy;
}

void Renderer::feedKey(int key, bool pressed) {
    if (key >= 0 && key < 512) {
        if (pressed && !impl->input.keys[key])
            impl->input.keyClicked[key] = true;
        impl->input.keys[key] = pressed;
    }
}

void Renderer::feedChar(unsigned int codepoint) {
    if (impl->input.charCount < 32)
        impl->input.charBuffer[impl->input.charCount++] = codepoint;
}

bool Renderer::isMouseOver(float x, float y, float w, float h) {
    return impl->input.mouseX >= x && impl->input.mouseX < x + w &&
           impl->input.mouseY >= y && impl->input.mouseY < y + h;
}

bool Renderer::isMouseClicked(float x, float y, float w, float h) {
    return isMouseOver(x, y, w, h) && impl->input.mouseClicked[0];
}

bool Renderer::isMouseDown(float x, float y, float w, float h) {
    return isMouseOver(x, y, w, h) && impl->input.mouseButtons[0];
}

bool Renderer::isKeyPressed(int key) {
    if (key < 0 || key >= 512) return false;
    return impl->input.keys[key];
}

bool Renderer::isKeyClicked(int key) {
    if (key < 0 || key >= 512) return false;
    return impl->input.keyClicked[key];
}

bool Renderer::isMouseDownAnywhere() { return impl->input.mouseButtons[0]; }
bool Renderer::isMouseClickedAnywhere() {
    return impl->input.mouseClicked[0];
}

float Renderer::mouseX() { return impl->input.mouseX; }
float Renderer::mouseY() { return impl->input.mouseY; }
float Renderer::scrollDeltaX() { return impl->input.scrollX; }
float Renderer::scrollDeltaY() { return impl->input.scrollY; }

int Renderer::pollChar() {
    if (impl->input.charCount <= 0) return -1;
    unsigned int c = impl->input.charBuffer[0];
    for (int i = 1; i < impl->input.charCount; i++)
        impl->input.charBuffer[i - 1] = impl->input.charBuffer[i];
    impl->input.charCount--;
    return static_cast<int>(c);
}

int Renderer::windowW() const { return impl->winW; }
int Renderer::windowH() const { return impl->winH; }

int Renderer::textureWidth(int id) const {
    if (id < 0 || id >= (int)impl->textures.size() || !impl->textures[id].active) return 0;
    return impl->textures[id].width;
}

int Renderer::textureHeight(int id) const {
    if (id < 0 || id >= (int)impl->textures.size() || !impl->textures[id].active) return 0;
    return impl->textures[id].height;
}

} // namespace vui