#include "vnc_pixel_codecs.h"

namespace VncPixelCodecs {

QRgb rgbFromPixelBytes(const uchar* bytes)
{
    // Byte order B,G,R,pad -- matches VncSessionBackend's negotiated
    // SetPixelFormat (32bpp little-endian, R at shift 16 / G at 8 / B at 0).
    return qRgb(bytes[2], bytes[1], bytes[0]);
}

int hextileFixedMetaByteCount(quint8 subencoding)
{
    int count = 0;
    if ((subencoding & HextileFlags::kBackgroundSpecified) != 0) {
        count += 4;
    }
    if ((subencoding & HextileFlags::kForegroundSpecified) != 0) {
        count += 4;
    }
    if ((subencoding & HextileFlags::kAnySubrects) != 0) {
        count += 1;
    }
    return count;
}

int decodeHextileFixedMeta(quint8 subencoding, const QByteArray& data, QRgb* background,
                           QRgb* foreground)
{
    int offset = 0;
    const auto* bytes = reinterpret_cast<const uchar*>(data.constData());

    if ((subencoding & HextileFlags::kBackgroundSpecified) != 0) {
        *background = rgbFromPixelBytes(bytes + offset);
        offset += 4;
    }
    if ((subencoding & HextileFlags::kForegroundSpecified) != 0) {
        *foreground = rgbFromPixelBytes(bytes + offset);
        offset += 4;
    }
    if ((subencoding & HextileFlags::kAnySubrects) != 0) {
        return static_cast<int>(static_cast<quint8>(data.at(offset)));
    }
    return 0;
}

int hextileSubrectByteCount(bool coloured, int subrectCount)
{
    return subrectCount * (coloured ? 6 : 2);
}

QVector<HextileSubrect> decodeHextileSubrects(bool coloured, int subrectCount,
                                              const QByteArray& data, QRgb foreground)
{
    QVector<HextileSubrect> subrects;
    subrects.reserve(subrectCount);

    const auto* bytes = reinterpret_cast<const uchar*>(data.constData());
    const int stride = coloured ? 6 : 2;
    for (int i = 0; i < subrectCount; ++i) {
        const uchar* entry = bytes + (i * stride);
        QRgb color = foreground;
        int fieldOffset = 0;
        if (coloured) {
            color = rgbFromPixelBytes(entry);
            fieldOffset = 4;
        }
        const uchar xy = entry[fieldOffset];
        const uchar wh = entry[fieldOffset + 1];
        // High nibble = x (or width-1), low nibble = y (or height-1).
        const int x = (xy >> 4) & 0x0F;
        const int y = xy & 0x0F;
        const int width = ((wh >> 4) & 0x0F) + 1;
        const int height = (wh & 0x0F) + 1;
        subrects.append(HextileSubrect{QRect(x, y, width, height), color});
    }
    return subrects;
}

int decodeZrleTile(const QByteArray& data, int offset, int tileWidth, int tileHeight,
                   QVector<QRgb>& pixels)
{
    const int pixelCount = tileWidth * tileHeight;
    const auto* bytes = reinterpret_cast<const uchar*>(data.constData());
    const int size = data.size();
    if (offset < 0 || offset >= size || pixelCount <= 0) {
        return -1;
    }

    int pos = offset;
    const quint8 subencoding = bytes[pos];
    ++pos;

    // CPIXEL is 3 bytes for our negotiated 32bpp/24-depth true-color
    // format -- the padding byte a full pixel would have is simply
    // omitted. rgbFromPixelBytes() already only reads the first 3 bytes it
    // is given (B,G,R order), so it doubles as the CPIXEL reader.
    auto readCpixel = [&](QRgb* out) -> bool {
        if (pos + 3 > size) {
            return false;
        }
        *out = rgbFromPixelBytes(bytes + pos);
        pos += 3;
        return true;
    };

    // Shared continuation-byte run-length reader for both RLE subencoding
    // families: keep summing bytes while they equal 255, add the final
    // (non-255) byte, and the true run length is that sum plus one.
    auto readRunLength = [&](int* out) -> bool {
        int total = 0;
        for (;;) {
            if (pos >= size) {
                return false;
            }
            const quint8 b = bytes[pos];
            ++pos;
            total += b;
            if (b != 255) {
                break;
            }
        }
        *out = total + 1;
        return true;
    };

    if (subencoding == 0) { // Raw
        pixels.reserve(pixels.size() + pixelCount);
        for (int i = 0; i < pixelCount; ++i) {
            QRgb color = 0;
            if (!readCpixel(&color)) {
                return -1;
            }
            pixels.append(color);
        }
        return pos - offset;
    }

    if (subencoding == 1) { // Solid
        QRgb color = 0;
        if (!readCpixel(&color)) {
            return -1;
        }
        pixels.reserve(pixels.size() + pixelCount);
        for (int i = 0; i < pixelCount; ++i) {
            pixels.append(color);
        }
        return pos - offset;
    }

    if (subencoding >= 2 && subencoding <= 16) { // Packed palette
        const int paletteSize = subencoding;
        QVector<QRgb> palette;
        palette.reserve(paletteSize);
        for (int i = 0; i < paletteSize; ++i) {
            QRgb color = 0;
            if (!readCpixel(&color)) {
                return -1;
            }
            palette.append(color);
        }

        int bitsPerPixel = 4;
        if (paletteSize == 2) {
            bitsPerPixel = 1;
        } else if (paletteSize <= 4) {
            bitsPerPixel = 2;
        }
        const int rowBytes = (tileWidth * bitsPerPixel + 7) / 8;

        pixels.reserve(pixels.size() + pixelCount);
        for (int y = 0; y < tileHeight; ++y) {
            if (pos + rowBytes > size) {
                return -1;
            }
            int bitPos = 0;
            for (int x = 0; x < tileWidth; ++x) {
                const int byteIndex = pos + (bitPos / 8);
                const int shift = 8 - (bitPos % 8) - bitsPerPixel;
                const int mask = (1 << bitsPerPixel) - 1;
                const int index = (bytes[byteIndex] >> shift) & mask;
                if (index >= palette.size()) {
                    return -1;
                }
                pixels.append(palette.at(index));
                bitPos += bitsPerPixel;
            }
            pos += rowBytes;
        }
        return pos - offset;
    }

    if (subencoding == 128) { // Plain RLE
        int produced = 0;
        while (produced < pixelCount) {
            QRgb color = 0;
            if (!readCpixel(&color)) {
                return -1;
            }
            int runLength = 0;
            if (!readRunLength(&runLength)) {
                return -1;
            }
            for (int i = 0; i < runLength && produced < pixelCount; ++i, ++produced) {
                pixels.append(color);
            }
        }
        return pos - offset;
    }

    if (subencoding >= 130) { // Palette RLE (129 is unused/invalid, falls through below)
        const int paletteSize = subencoding - 128;
        QVector<QRgb> palette;
        palette.reserve(paletteSize);
        for (int i = 0; i < paletteSize; ++i) {
            QRgb color = 0;
            if (!readCpixel(&color)) {
                return -1;
            }
            palette.append(color);
        }

        int produced = 0;
        while (produced < pixelCount) {
            if (pos >= size) {
                return -1;
            }
            const quint8 indexByte = bytes[pos];
            ++pos;

            int index = indexByte;
            int runLength = 1;
            if (indexByte >= 128) {
                index = indexByte - 128;
                if (!readRunLength(&runLength)) {
                    return -1;
                }
            }
            if (index >= palette.size()) {
                return -1;
            }
            const QRgb color = palette.at(index);
            for (int i = 0; i < runLength && produced < pixelCount; ++i, ++produced) {
                pixels.append(color);
            }
        }
        return pos - offset;
    }

    // Subencodings 17-127 and 129 are not defined by RFC 6143.
    return -1;
}

QVector<QRgb> decodeTightCopyFilter(const QByteArray& data, int width, int height)
{
    const int pixelCount = width * height;
    QVector<QRgb> pixels;
    if (pixelCount <= 0 || data.size() < pixelCount * 3) {
        return pixels;
    }
    pixels.reserve(pixelCount);
    const auto* bytes = reinterpret_cast<const uchar*>(data.constData());
    for (int i = 0; i < pixelCount; ++i) {
        pixels.append(rgbFromPixelBytes(bytes + (i * 3)));
    }
    return pixels;
}

QVector<QRgb> decodeTightPaletteFilter(const QByteArray& data, int width, int height)
{
    QVector<QRgb> pixels;
    const int pixelCount = width * height;
    if (pixelCount <= 0 || data.isEmpty()) {
        return pixels;
    }

    const auto* bytes = reinterpret_cast<const uchar*>(data.constData());
    const int size = data.size();
    int pos = 0;
    const int paletteSize = static_cast<int>(bytes[pos]) + 1; // 1-256 colors
    ++pos;
    if (pos + (paletteSize * 3) > size) {
        return pixels;
    }

    QVector<QRgb> palette;
    palette.reserve(paletteSize);
    for (int i = 0; i < paletteSize; ++i) {
        palette.append(rgbFromPixelBytes(bytes + pos));
        pos += 3;
    }

    int bitsPerPixel = 8;
    if (paletteSize <= 2) {
        bitsPerPixel = 1;
    } else if (paletteSize <= 4) {
        bitsPerPixel = 2;
    } else if (paletteSize <= 16) {
        bitsPerPixel = 4;
    }

    pixels.reserve(pixelCount);
    int bitPos = 0;
    for (int i = 0; i < pixelCount; ++i) {
        const int byteIndex = pos + (bitPos / 8);
        if (byteIndex >= size) {
            return QVector<QRgb>();
        }
        const int shift = 8 - (bitPos % 8) - bitsPerPixel;
        const int mask = (1 << bitsPerPixel) - 1;
        const int index = (bytes[byteIndex] >> shift) & mask;
        if (index >= palette.size()) {
            return QVector<QRgb>();
        }
        pixels.append(palette.at(index));
        bitPos += bitsPerPixel;
    }
    return pixels;
}

QVector<QRgb> decodeTightGradientFilter(const QByteArray& data, int width, int height)
{
    const int pixelCount = width * height;
    QVector<QRgb> pixels;
    if (pixelCount <= 0 || data.size() < pixelCount * 3) {
        return pixels;
    }

    const auto* bytes = reinterpret_cast<const uchar*>(data.constData());
    auto predict = [](int left, int up, int upLeft) {
        return qBound(0, left + up - upLeft, 255);
    };

    QVector<int> rChan(pixelCount);
    QVector<int> gChan(pixelCount);
    QVector<int> bChan(pixelCount);
    for (int y = 0; y < height; ++y) {
        for (int x = 0; x < width; ++x) {
            const int idx = (y * width) + x;
            const uchar* px = bytes + (idx * 3);
            // TPIXEL order matches rgbFromPixelBytes: B,G,R.
            const int deltaB = px[0];
            const int deltaG = px[1];
            const int deltaR = px[2];

            const int leftR = (x > 0) ? rChan[idx - 1] : 0;
            const int leftG = (x > 0) ? gChan[idx - 1] : 0;
            const int leftB = (x > 0) ? bChan[idx - 1] : 0;
            const int upR = (y > 0) ? rChan[idx - width] : 0;
            const int upG = (y > 0) ? gChan[idx - width] : 0;
            const int upB = (y > 0) ? bChan[idx - width] : 0;
            const int upLeftR = (x > 0 && y > 0) ? rChan[idx - width - 1] : 0;
            const int upLeftG = (x > 0 && y > 0) ? gChan[idx - width - 1] : 0;
            const int upLeftB = (x > 0 && y > 0) ? bChan[idx - width - 1] : 0;

            rChan[idx] = (predict(leftR, upR, upLeftR) + deltaR) & 0xFF;
            gChan[idx] = (predict(leftG, upG, upLeftG) + deltaG) & 0xFF;
            bChan[idx] = (predict(leftB, upB, upLeftB) + deltaB) & 0xFF;
        }
    }

    pixels.reserve(pixelCount);
    for (int i = 0; i < pixelCount; ++i) {
        pixels.append(qRgb(rChan.at(i), gChan.at(i), bChan.at(i)));
    }
    return pixels;
}

}