/*
    * framebuffer.hpp
    * MontaukOS double-buffered framebuffer abstraction
    * Copyright (c) 2025 Daniel Hammer
*/

#pragma once
#include <cstdint>
#include <montauk/syscall.h>
#include <montauk/string.h>
#include "gui/gui.hpp"

namespace gui {

class Framebuffer {
    uint32_t* hw_fb;
    uint32_t* back_buf;
    int fb_width;
    int fb_height;
    int fb_pitch; // in bytes

    static inline void fill_pixels(uint32_t* dst, int count, uint32_t pixel) {
        if (!dst || count <= 0) return;

        uint64_t pixel64 = ((uint64_t)pixel << 32) | pixel;
        if (((uint64_t)dst & 4) && count > 0) {
            *dst++ = pixel;
            count--;
        }

        uint64_t* dst64 = (uint64_t*)dst;
        int pairs = count / 2;
        for (int i = 0; i < pairs; i++) {
            dst64[i] = pixel64;
        }

        if (count & 1) {
            ((uint32_t*)(dst64 + pairs))[0] = pixel;
        }
    }

    static inline uint32_t blend_pixel(uint32_t dst, uint32_t src, uint32_t alpha) {
        uint32_t inv_a = 255 - alpha;
        uint32_t sr = (src >> 16) & 0xFF;
        uint32_t sg = (src >> 8) & 0xFF;
        uint32_t sb = src & 0xFF;
        uint32_t dr = (dst >> 16) & 0xFF;
        uint32_t dg = (dst >> 8) & 0xFF;
        uint32_t db = dst & 0xFF;

        uint32_t rr = alpha * sr + inv_a * dr;
        uint32_t gg = alpha * sg + inv_a * dg;
        uint32_t bb = alpha * sb + inv_a * db;

        rr = (rr + 1 + (rr >> 8)) >> 8;
        gg = (gg + 1 + (gg >> 8)) >> 8;
        bb = (bb + 1 + (bb >> 8)) >> 8;

        return 0xFF000000 | (rr << 16) | (gg << 8) | bb;
    }

public:
    Framebuffer() : hw_fb(nullptr), back_buf(nullptr), fb_width(0), fb_height(0), fb_pitch(0) {
        montauk::abi::FbInfo info;
        montauk::fb_info(&info);

        fb_width  = (int)info.width;
        fb_height = (int)info.height;
        fb_pitch  = (int)info.pitch;

        hw_fb = (uint32_t*)montauk::fb_map();
        back_buf = (uint32_t*)montauk::alloc((uint64_t)fb_height * fb_pitch);
    }

    int width() const { return fb_width; }
    int height() const { return fb_height; }
    int pitch() const { return fb_pitch; }

    uint32_t* buffer() { return back_buf; }

    inline void put_pixel(int x, int y, Color c) {
        if (x < 0 || x >= fb_width || y < 0 || y >= fb_height) return;
        uint32_t* row = (uint32_t*)((uint8_t*)back_buf + y * fb_pitch);
        row[x] = c.to_pixel();
    }

    inline void put_pixel_alpha(int x, int y, Color c) {
        if (x < 0 || x >= fb_width || y < 0 || y >= fb_height) return;
        if (c.a == 0) return;
        if (c.a == 255) {
            put_pixel(x, y, c);
            return;
        }

        uint32_t* row = (uint32_t*)((uint8_t*)back_buf + y * fb_pitch);
        uint32_t dst = row[x];

        uint8_t dr = (dst >> 16) & 0xFF;
        uint8_t dg = (dst >> 8) & 0xFF;
        uint8_t db = dst & 0xFF;

        uint32_t a = c.a;
        uint32_t inv_a = 255 - a;

        // Fast alpha blend: out = (src * alpha + dst * (255 - alpha) + 128) / 255
        // Approximation: (x + 1 + (x >> 8)) >> 8 for division by 255
        uint32_t rr = a * c.r + inv_a * dr;
        uint32_t gg = a * c.g + inv_a * dg;
        uint32_t bb = a * c.b + inv_a * db;

        rr = (rr + 1 + (rr >> 8)) >> 8;
        gg = (gg + 1 + (gg >> 8)) >> 8;
        bb = (bb + 1 + (bb >> 8)) >> 8;

        row[x] = (0xFF000000) | (rr << 16) | (gg << 8) | bb;
    }

    inline void fill_rect(int x, int y, int w, int h, Color c) {
        if (!back_buf) return;

        // Clip to screen bounds
        int x0 = x < 0 ? 0 : x;
        int y0 = y < 0 ? 0 : y;
        int x1 = (x + w) > fb_width ? fb_width : (x + w);
        int y1 = (y + h) > fb_height ? fb_height : (y + h);

        if (x0 >= x1 || y0 >= y1) return;

        uint32_t pixel = c.to_pixel();
        int clipped_w = x1 - x0;

        for (int row = y0; row < y1; row++) {
            uint32_t* dst = (uint32_t*)((uint8_t*)back_buf + row * fb_pitch) + x0;
            fill_pixels(dst, clipped_w, pixel);
        }
    }

    inline void fill_rect_alpha(int x, int y, int w, int h, Color c) {
        if (c.a == 0) return;
        if (c.a == 255) {
            fill_rect(x, y, w, h, c);
            return;
        }

        int x0 = x < 0 ? 0 : x;
        int y0 = y < 0 ? 0 : y;
        int x1 = (x + w) > fb_width ? fb_width : (x + w);
        int y1 = (y + h) > fb_height ? fb_height : (y + h);

        if (x0 >= x1 || y0 >= y1) return;

        uint32_t a = c.a;
        uint32_t inv_a = 255 - a;
        uint32_t src_r = a * c.r;
        uint32_t src_g = a * c.g;
        uint32_t src_b = a * c.b;

        for (int row = y0; row < y1; row++) {
            uint32_t* dst = (uint32_t*)((uint8_t*)back_buf + row * fb_pitch) + x0;
            for (int col = 0; col < x1 - x0; col++) {
                uint32_t d = dst[col];
                uint32_t dr = (d >> 16) & 0xFF;
                uint32_t dg = (d >> 8) & 0xFF;
                uint32_t db = d & 0xFF;

                uint32_t rr = src_r + inv_a * dr;
                uint32_t gg = src_g + inv_a * dg;
                uint32_t bb = src_b + inv_a * db;

                rr = (rr + 1 + (rr >> 8)) >> 8;
                gg = (gg + 1 + (gg >> 8)) >> 8;
                bb = (bb + 1 + (bb >> 8)) >> 8;

                dst[col] = (0xFF000000) | (rr << 16) | (gg << 8) | bb;
            }
        }
    }

    inline void blit(int x, int y, int w, int h, const uint32_t* pixels) {
        if (!back_buf || !pixels || w <= 0 || h <= 0) return;

        int src_x = 0;
        int src_y = 0;
        int dst_x = x;
        int dst_y = y;
        int copy_w = w;
        int copy_h = h;

        if (dst_x < 0) { src_x = -dst_x; copy_w += dst_x; dst_x = 0; }
        if (dst_y < 0) { src_y = -dst_y; copy_h += dst_y; dst_y = 0; }
        if (dst_x + copy_w > fb_width)  copy_w = fb_width - dst_x;
        if (dst_y + copy_h > fb_height) copy_h = fb_height - dst_y;
        if (copy_w <= 0 || copy_h <= 0) return;

        uint64_t row_bytes = (uint64_t)copy_w * sizeof(uint32_t);
        for (int row = 0; row < copy_h; row++) {
            uint32_t* dst = (uint32_t*)((uint8_t*)back_buf + (dst_y + row) * fb_pitch) + dst_x;
            const uint32_t* src = pixels + (src_y + row) * w + src_x;
            montauk::memcpy(dst, src, row_bytes);
        }
    }

    inline void blit_alpha(int x, int y, int w, int h, const uint32_t* pixels) {
        if (!back_buf || !pixels || w <= 0 || h <= 0) return;

        int src_x = 0;
        int src_y = 0;
        int dst_x = x;
        int dst_y = y;
        int copy_w = w;
        int copy_h = h;

        if (dst_x < 0) { src_x = -dst_x; copy_w += dst_x; dst_x = 0; }
        if (dst_y < 0) { src_y = -dst_y; copy_h += dst_y; dst_y = 0; }
        if (dst_x + copy_w > fb_width)  copy_w = fb_width - dst_x;
        if (dst_y + copy_h > fb_height) copy_h = fb_height - dst_y;
        if (copy_w <= 0 || copy_h <= 0) return;

        for (int row = 0; row < copy_h; row++) {
            uint32_t* dst = (uint32_t*)((uint8_t*)back_buf + (dst_y + row) * fb_pitch) + dst_x;
            const uint32_t* src = pixels + (src_y + row) * w + src_x;
            for (int col = 0; col < copy_w; col++) {
                uint32_t s = src[col];
                uint8_t sa = (s >> 24) & 0xFF;
                if (sa == 0) continue;
                if (sa == 255) {
                    dst[col] = s;
                    continue;
                }

                dst[col] = blend_pixel(dst[col], s, sa);
            }
        }
    }

    inline void copy_from(const uint32_t* pixels, int src_pitch_bytes) {
        if (!back_buf || !pixels || src_pitch_bytes <= 0) return;
        uint64_t row_bytes = (uint64_t)fb_width * sizeof(uint32_t);
        for (int row = 0; row < fb_height; row++) {
            uint32_t* dst = (uint32_t*)((uint8_t*)back_buf + row * fb_pitch);
            const uint32_t* src = (const uint32_t*)((const uint8_t*)pixels + row * src_pitch_bytes);
            montauk::memcpy(dst, src, row_bytes);
        }
    }

    inline void clear(Color c) {
        fill_rect(0, 0, fb_width, fb_height, c);
    }

    inline void flip() {
        if (!hw_fb || !back_buf) return;

        // Copy back buffer to hardware framebuffer, row by row (pitch may differ)
        uint64_t row_bytes = (uint64_t)fb_width * sizeof(uint32_t);
        for (int y = 0; y < fb_height; y++) {
            uint32_t* src = (uint32_t*)((uint8_t*)back_buf + y * fb_pitch);
            uint32_t* dst = (uint32_t*)((uint8_t*)hw_fb + y * fb_pitch);
            montauk::memcpy(dst, src, row_bytes);
        }
    }
};

} // namespace gui