#include "core/Image.h"

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

#define STB_IMAGE_IMPLEMENTATION
#define STBI_NO_PSD
#define STBI_NO_TGA
#define STBI_NO_HDR
#define STBI_NO_PIC
#define STBI_NO_PNM
// We deliberately do NOT disable STBI_NO_BMP: stb_image's BMP path is our
// fallback for 24/32bpp BMPs that aren't indexed. Indexed BMPs are read by
// our own reader below instead, before stb_image ever sees the file.
#include "stb_image.h"

namespace bmp2tile
{

namespace
{

// ---- tiny little-endian byte readers -----------------------------------

uint16_t ReadU16(const uint8_t* p)
{
    return static_cast<uint16_t>(p[0] | (p[1] << 8));
}

uint32_t ReadU32(const uint8_t* p)
{
    return static_cast<uint32_t>(p[0]) | (static_cast<uint32_t>(p[1]) << 8) |
           (static_cast<uint32_t>(p[2]) << 16) | (static_cast<uint32_t>(p[3]) << 24);
}

int32_t ReadI32(const uint8_t* p)
{
    return static_cast<int32_t>(ReadU32(p));
}

// ---- indexed BMP reader --------------------------------------------------
//
// Handles the common case: BITMAPFILEHEADER (14 bytes) + a 40-byte
// BITMAPINFOHEADER, uncompressed (BI_RGB), 1/4/8 bits per pixel, with a
// palette of BGR(A) quads. This covers what every mainstream paint program
// produces when you "save as 16/256-colour bitmap", which is exactly the
// workflow the original tool's README describes.
//
// Returns false (rather than throwing) if the file isn't a BMP at all, so
// the caller can fall through to the stb_image path.
bool TryReadIndexedBmp(const std::vector<uint8_t>& data, IndexedImage& out)
{
    if (data.size() < 14 + 40 || data[0] != 'B' || data[1] != 'M')
    {
        return false;  // not a BMP; let stb_image have a go
    }

    const uint32_t pixelDataOffset = ReadU32(&data[10]);
    const uint32_t infoHeaderSize = ReadU32(&data[14]);
    if (infoHeaderSize < 40)
    {
        throw ImageLoadError("Unsupported BMP variant (header size " + std::to_string(infoHeaderSize) +
                              " bytes) -- re-save as a standard Windows BMP.");
    }

    const int32_t rawWidth = ReadI32(&data[18]);
    const int32_t rawHeight = ReadI32(&data[22]);
    const uint16_t bitCount = ReadU16(&data[28]);
    const uint32_t compression = ReadU32(&data[30]);
    uint32_t colorsUsed = ReadU32(&data[46]);

    if (bitCount != 1 && bitCount != 4 && bitCount != 8)
    {
        // 16/24/32bpp BMP: not indexed, hand off to stb_image.
        return false;
    }
    if (compression != 0 /* BI_RGB */)
    {
        throw ImageLoadError("Compressed/RLE BMPs aren't supported -- re-save uncompressed.");
    }

    const int width = rawWidth;
    const bool bottomUp = rawHeight > 0;
    const int height = bottomUp ? rawHeight : -rawHeight;
    if (width <= 0 || height <= 0)
    {
        throw ImageLoadError("BMP reports a non-positive width/height.");
    }

    const uint32_t maxColors = 1u << bitCount;
    if (colorsUsed == 0 || colorsUsed > maxColors)
    {
        colorsUsed = maxColors;
    }

    const size_t paletteOffset = 14 + infoHeaderSize;
    if (paletteOffset + static_cast<size_t>(colorsUsed) * 4 > data.size())
    {
        throw ImageLoadError("BMP palette is truncated.");
    }

    out.Palette.resize(colorsUsed);
    for (uint32_t i = 0; i < colorsUsed; ++i)
    {
        const uint8_t* entry = &data[paletteOffset + static_cast<size_t>(i) * 4];
        // Palette entries are stored B, G, R, reserved.
        out.Palette[i] = Color{entry[2], entry[1], entry[0]};
    }

    out.Width = width;
    out.Height = height;
    out.Pixels.assign(static_cast<size_t>(width) * height, 0);

    const size_t rowBytes = ((static_cast<size_t>(width) * bitCount + 31) / 32) * 4;
    if (pixelDataOffset + rowBytes * height > data.size())
    {
        throw ImageLoadError("BMP pixel data is truncated.");
    }

    for (int fileRow = 0; fileRow < height; ++fileRow)
    {
        // BMP rows are bottom-up by default (positive height); a negative
        // height means top-down. Either way we write into `out` top-down.
        const int destY = bottomUp ? (height - 1 - fileRow) : fileRow;
        const uint8_t* row = &data[pixelDataOffset + static_cast<size_t>(fileRow) * rowBytes];

        for (int x = 0; x < width; ++x)
        {
            uint8_t index = 0;
            switch (bitCount)
            {
                case 1:
                    index = (row[x / 8] >> (7 - (x % 8))) & 0x1;
                    break;
                case 4:
                    index = (x % 2 == 0) ? (row[x / 2] >> 4) : (row[x / 2] & 0x0F);
                    break;
                case 8:
                    index = row[x];
                    break;
            }
            if (index >= out.Palette.size())
            {
                // Be lenient rather than throwing: clamp, since some tools
                // write a slightly-too-small palette table for images that
                // don't actually use every entry.
                index = static_cast<uint8_t>(out.Palette.size() - 1);
            }
            out.Pixels[static_cast<size_t>(destY) * width + x] = index;
        }
    }

    return true;
}

// ---- stb_image fallback: decode to RGBA, then quantize by unique colour --

IndexedImage LoadViaStbAndQuantize(const std::string& path, int maxColors)
{
    int width = 0, height = 0, channelsInFile = 0;
    stbi_uc* rgba = stbi_load(path.c_str(), &width, &height, &channelsInFile, 4);
    if (!rgba)
    {
        throw ImageLoadError("Couldn't decode image: " + std::string(stbi_failure_reason()));
    }

    IndexedImage out;
    out.Width = width;
    out.Height = height;
    out.Pixels.resize(static_cast<size_t>(width) * height);

    // Maps a packed 0x00RRGGBB colour to its palette index, preserving the
    // order colours are first encountered in (matches the "the palette
    // reflects what's in the image" behaviour of the original tool, even
    // though the *specific order* won't match an embedded palette table).
    std::unordered_map<uint32_t, uint8_t> colorToIndex;
    colorToIndex.reserve(static_cast<size_t>(maxColors) * 2);

    for (int i = 0; i < width * height; ++i)
    {
        const uint8_t r = rgba[i * 4 + 0];
        const uint8_t g = rgba[i * 4 + 1];
        const uint8_t b = rgba[i * 4 + 2];
        const uint32_t packed = (static_cast<uint32_t>(r) << 16) | (static_cast<uint32_t>(g) << 8) | b;

        auto it = colorToIndex.find(packed);
        if (it != colorToIndex.end())
        {
            out.Pixels[i] = it->second;
            continue;
        }

        if (static_cast<int>(out.Palette.size()) >= maxColors)
        {
            stbi_image_free(rgba);
            throw ImageLoadError("Image uses more than " + std::to_string(maxColors) +
                                  " colours -- reduce it to an indexed palette first "
                                  "(this is also why indexed BMP is the recommended input format).");
        }

        const uint8_t newIndex = static_cast<uint8_t>(out.Palette.size());
        out.Palette.push_back(Color{r, g, b});
        colorToIndex[packed] = newIndex;
        out.Pixels[i] = newIndex;
    }

    stbi_image_free(rgba);
    return out;
}

std::vector<uint8_t> ReadWholeFile(const std::string& path)
{
    std::ifstream f(path, std::ios::binary | std::ios::ate);
    if (!f)
    {
        throw ImageLoadError("Couldn't open file: " + path);
    }
    const std::streamsize size = f.tellg();
    f.seekg(0);
    std::vector<uint8_t> data(static_cast<size_t>(size));
    if (!f.read(reinterpret_cast<char*>(data.data()), size))
    {
        throw ImageLoadError("Couldn't read file: " + path);
    }
    return data;
}

}  // namespace

IndexedImage LoadImage(const std::string& path, int maxColors)
{
    // PCX is explicitly out of scope for this port.
    if (path.size() >= 4)
    {
        std::string ext = path.substr(path.size() - 4);
        std::transform(ext.begin(), ext.end(), ext.begin(), ::tolower);
        if (ext == ".pcx")
        {
            throw ImageLoadError("PCX is not supported by this build -- re-save as BMP, PNG or GIF.");
        }
    }

    const std::vector<uint8_t> data = ReadWholeFile(path);

    IndexedImage image;
    if (TryReadIndexedBmp(data, image))
    {
        return image;
    }

    return LoadViaStbAndQuantize(path, maxColors);
}

}  // namespace bmp2tile
