#include "blp_loader.h"

#include <godot_cpp/core/class_db.hpp>
#include <godot_cpp/variant/utility_functions.hpp>

#include <fstream>
#include <cstring>
#include <algorithm>

using namespace godot;

// ─────────────────────────────────────────────────────────────────────────────
// BLP2 binary structures
// ─────────────────────────────────────────────────────────────────────────────
#pragma pack(push, 1)
struct BLP2Header {
    char     magic[4];          // "BLP2"
    uint32_t type;              // 0=JPEG 1=palette 2=DXT
    uint8_t  encoding;          // 1=raw(palette) 2=DXT 3=raw(BGRA)
    uint8_t  alphaDepth;        // 0,1,4,8
    uint8_t  alphaEncoding;     // 0=DXT1 1=DXT3 7=DXT5
    uint8_t  hasMipmaps;
    uint32_t width;
    uint32_t height;
    uint32_t mipmapOffset[16];
    uint32_t mipmapSize[16];
};
#pragma pack(pop)

// ─────────────────────────────────────────────────────────────────────────────
// DXT colour block helpers
// ─────────────────────────────────────────────────────────────────────────────
static void unpack_565(uint16_t c, uint8_t &r, uint8_t &g, uint8_t &b) {
    r = (c >> 11) & 0x1F; r = (r << 3) | (r >> 2);
    g = (c >>  5) & 0x3F; g = (g << 2) | (g >> 4);
    b =  c        & 0x1F; b = (b << 3) | (b >> 2);
}

// Decode a 4×4 DXT1 colour block into dst (RGBA row-major, y=block_y*4+row)
static void decode_dxt1_block(const uint8_t *src,
                               uint8_t *dst, int pitch, bool has_alpha)
{
    uint16_t c0 = src[0] | (src[1] << 8);
    uint16_t c1 = src[2] | (src[3] << 8);
    uint8_t r[4], g[4], b[4], a[4];
    unpack_565(c0, r[0], g[0], b[0]); a[0] = 255;
    unpack_565(c1, r[1], g[1], b[1]); a[1] = 255;

    if (c0 > c1 || !has_alpha) {
        r[2] = (2*r[0]+r[1])/3; g[2] = (2*g[0]+g[1])/3; b[2] = (2*b[0]+b[1])/3; a[2] = 255;
        r[3] = (r[0]+2*r[1])/3; g[3] = (g[0]+2*g[1])/3; b[3] = (b[0]+2*b[1])/3; a[3] = 255;
    } else {
        r[2] = (r[0]+r[1])/2;   g[2] = (g[0]+g[1])/2;   b[2] = (b[0]+b[1])/2;   a[2] = 255;
        r[3] = 0; g[3] = 0; b[3] = 0; a[3] = 0;
    }

    uint32_t indices = src[4] | (src[5]<<8) | (src[6]<<16) | (src[7]<<24);
    for (int row = 0; row < 4; ++row) {
        for (int col = 0; col < 4; ++col) {
            int i = (indices >> (2*(row*4+col))) & 3;
            uint8_t *p = dst + row * pitch + col * 4;
            p[0] = r[i]; p[1] = g[i]; p[2] = b[i]; p[3] = a[i];
        }
    }
}

// DXT3 explicit alpha block (16 bytes: 8 alpha + 8 colour)
static void decode_dxt3_block(const uint8_t *src, uint8_t *dst, int pitch) {
    uint8_t alpha[16];
    for (int i = 0; i < 8; ++i) {
        alpha[i*2]   = (src[i] & 0x0F) * 17;
        alpha[i*2+1] = (src[i] >> 4)   * 17;
    }
    decode_dxt1_block(src + 8, dst, pitch, false);
    for (int row = 0; row < 4; ++row)
        for (int col = 0; col < 4; ++col)
            dst[row * pitch + col * 4 + 3] = alpha[row*4+col];
}

// DXT5 interpolated alpha block (16 bytes: 8 alpha + 8 colour)
static void decode_dxt5_block(const uint8_t *src, uint8_t *dst, int pitch) {
    uint8_t a0 = src[0], a1 = src[1];
    uint8_t atable[8];
    atable[0] = a0; atable[1] = a1;
    if (a0 > a1) {
        for (int i = 2; i < 8; ++i)
            atable[i] = ((8-i)*a0 + (i-1)*a1) / 7;
    } else {
        for (int i = 2; i < 6; ++i)
            atable[i] = ((6-i)*a0 + (i-1)*a1) / 5;
        atable[6] = 0; atable[7] = 255;
    }
    uint64_t abits = 0;
    for (int i = 0; i < 6; ++i) abits |= ((uint64_t)src[2+i]) << (i*8);

    decode_dxt1_block(src + 8, dst, pitch, false);
    for (int row = 0; row < 4; ++row)
        for (int col = 0; col < 4; ++col) {
            int bit = row*4+col;
            int idx = (abits >> (bit*3)) & 7;
            dst[row*pitch + col*4 + 3] = atable[idx];
        }
}

// ─────────────────────────────────────────────────────────────────────────────
// Full mip-level decoders
// ─────────────────────────────────────────────────────────────────────────────
void BLPLoader::decode_dxt1(const uint8_t *src, int w, int h,
                             std::vector<uint8_t> &dst, bool has_alpha)
{
    int bw = (w + 3) / 4, bh = (h + 3) / 4;
    dst.assign(w * h * 4, 0);
    for (int by = 0; by < bh; ++by)
        for (int bx = 0; bx < bw; ++bx, src += 8) {
            uint8_t *d = dst.data() + (by*4*w + bx*4) * 4;
            decode_dxt1_block(src, d, w*4, has_alpha);
        }
}

void BLPLoader::decode_dxt3(const uint8_t *src, int w, int h,
                             std::vector<uint8_t> &dst)
{
    int bw = (w+3)/4, bh = (h+3)/4;
    dst.assign(w * h * 4, 0);
    for (int by = 0; by < bh; ++by)
        for (int bx = 0; bx < bw; ++bx, src += 16) {
            uint8_t *d = dst.data() + (by*4*w + bx*4)*4;
            decode_dxt3_block(src, d, w*4);
        }
}

void BLPLoader::decode_dxt5(const uint8_t *src, int w, int h,
                             std::vector<uint8_t> &dst)
{
    int bw = (w+3)/4, bh = (h+3)/4;
    dst.assign(w * h * 4, 0);
    for (int by = 0; by < bh; ++by)
        for (int bx = 0; bx < bw; ++bx, src += 16) {
            uint8_t *d = dst.data() + (by*4*w + bx*4)*4;
            decode_dxt5_block(src, d, w*4);
        }
}

// ─────────────────────────────────────────────────────────────────────────────
// Main parser
// ─────────────────────────────────────────────────────────────────────────────
Ref<Image> BLPLoader::parse(const uint8_t *data, size_t len) {
    if (len < sizeof(BLP2Header)) return Ref<Image>();

    const BLP2Header &hdr = *reinterpret_cast<const BLP2Header *>(data);
    if (std::memcmp(hdr.magic, "BLP2", 4) != 0) {
        // Try BLP1 magic
        if (std::memcmp(hdr.magic, "BLP1", 4) != 0) return Ref<Image>();
    }

    uint32_t w = hdr.width, h = hdr.height;
    if (w == 0 || h == 0) return Ref<Image>();

    // Use mip level 0 only
    uint32_t mip_ofs  = hdr.mipmapOffset[0];
    uint32_t mip_size = hdr.mipmapSize[0];
    if (mip_ofs + mip_size > len) return Ref<Image>();
    const uint8_t *mip = data + mip_ofs;

    std::vector<uint8_t> rgba;
    Image::Format fmt = Image::FORMAT_RGBA8;

    if (hdr.encoding == 2) {
        // DXT compressed
        bool has_alpha = hdr.alphaDepth > 0;
        switch (hdr.alphaEncoding) {
            case 0: decode_dxt1(mip, w, h, rgba, has_alpha); break;
            case 1: decode_dxt3(mip, w, h, rgba); break;
            case 7: decode_dxt5(mip, w, h, rgba); break;
            default: decode_dxt1(mip, w, h, rgba, has_alpha); break;
        }

    } else if (hdr.encoding == 1) {
        // Palettized: 256×4 BGRA palette at offset 0xA4, then index data
        const size_t palette_ofs = sizeof(BLP2Header);
        if (palette_ofs + 256*4 > len) return Ref<Image>();
        const uint8_t *palette = data + palette_ofs;
        const uint8_t *alpha = mip + (size_t)w * (size_t)h;

        rgba.resize(w * h * 4);
        for (uint32_t i = 0; i < w * h && i < mip_size; ++i) {
            uint8_t idx = mip[i];
            rgba[i*4+0] = palette[idx*4+2]; // R (from BGR)
            rgba[i*4+1] = palette[idx*4+1]; // G
            rgba[i*4+2] = palette[idx*4+0]; // B
            uint8_t a = 255;
            if (hdr.alphaDepth == 8) {
                size_t alpha_pos = (size_t)w * (size_t)h + i;
                a = alpha_pos < mip_size ? mip[alpha_pos] : 255;
            } else if (hdr.alphaDepth == 4) {
                size_t byte_pos = i >> 1;
                if ((size_t)w * (size_t)h + byte_pos < mip_size) {
                    uint8_t nibble = (i & 1) ? (alpha[byte_pos] >> 4) : (alpha[byte_pos] & 0x0F);
                    a = (uint8_t)(nibble * 17);
                }
            } else if (hdr.alphaDepth == 1) {
                size_t byte_pos = i >> 3;
                if ((size_t)w * (size_t)h + byte_pos < mip_size) {
                    a = (alpha[byte_pos] & (1 << (i & 7))) ? 255 : 0;
                }
            }
            rgba[i*4+3] = a;
        }

    } else if (hdr.encoding == 3) {
        // Uncompressed BGRA
        rgba.resize(w * h * 4);
        for (uint32_t i = 0; i < w * h * 4 && i + 3 < mip_size; i += 4) {
            rgba[i+0] = mip[i+2]; // R
            rgba[i+1] = mip[i+1]; // G
            rgba[i+2] = mip[i+0]; // B
            rgba[i+3] = mip[i+3]; // A
        }
    } else {
        return Ref<Image>();
    }

    if (rgba.empty()) return Ref<Image>();

    PackedByteArray pba;
    pba.resize((int)rgba.size());
    std::memcpy(pba.ptrw(), rgba.data(), rgba.size());

    return Image::create_from_data(w, h, false, fmt, pba);
}

// ─────────────────────────────────────────────────────────────────────────────
// Public API
// ─────────────────────────────────────────────────────────────────────────────
Ref<Image> BLPLoader::load_image(const String &path) {
    auto buf = read_file(to_std(path));
    if (buf.empty()) {
        UtilityFunctions::push_error("BLPLoader: cannot read ", path);
        return Ref<Image>();
    }
    return parse(buf.data(), buf.size());
}

Ref<Image> BLPLoader::load_image_from_bytes(const PackedByteArray &bytes) {
    if (bytes.is_empty()) return Ref<Image>();
    return parse(bytes.ptr(), bytes.size());
}

std::vector<uint8_t> BLPLoader::read_file(const std::string &path) {
    std::ifstream f(path, std::ios::binary | std::ios::ate);
    if (!f) return {};
    auto sz = f.tellg(); f.seekg(0);
    std::vector<uint8_t> buf(sz);
    f.read(reinterpret_cast<char *>(buf.data()), sz);
    return buf;
}

std::string BLPLoader::to_std(const String &s) {
    return std::string(s.utf8().get_data());
}

void BLPLoader::_bind_methods() {
    ClassDB::bind_method(D_METHOD("load_image",            "path"),  &BLPLoader::load_image);
    ClassDB::bind_method(D_METHOD("load_image_from_bytes", "bytes"), &BLPLoader::load_image_from_bytes);
}
