#include "PixelFormat.hpp"

#include <cstdint>
#include <cstdio>
#include <cstring>
#include <cmath>
#include <vector>

namespace
{
    struct YCbCr10
    {
        uint32_t y;
        uint32_t cb;
        uint32_t cr;
    };

    struct BarSpec
    {
        const char* name;
        RgbColor rgb75;
    };

    const BarSpec bars[] = {
        {"White",   {191, 191, 191}},
        {"Yellow",  {191, 191,   0}},
        {"Cyan",    {  0, 191, 191}},
        {"Green",   {  0, 191,   0}},
        {"Magenta", {191,   0, 191}},
        {"Red",     {191,   0,   0}},
        {"Blue",    {  0,   0, 191}},
        {"Black",   {  0,   0,   0}},
    };

    YCbCr10 rgbToYCbCr10(RgbColor rgb)
    {
        double r = rgb.r / 255.0;
        double g = rgb.g / 255.0;
        double b = rgb.b / 255.0;

        double y  =  0.2126 * r + 0.7152 * g + 0.0722 * b;
        double cb = -0.1146 * r - 0.3854 * g + 0.5000 * b;
        double cr =  0.5000 * r - 0.4542 * g - 0.0458 * b;

        uint32_t y10  = static_cast<uint32_t>(std::round(y * 876.0 + 64.0));
        uint32_t cb10 = static_cast<uint32_t>(std::round(cb * 896.0 + 512.0));
        uint32_t cr10 = static_cast<uint32_t>(std::round(cr * 896.0 + 512.0));

        return {y10, cb10, cr10};
    }

    void encodeV210Line(
        const YCbCr10* pixels,
        uint32_t width,
        uint8_t* out)
    {
        uint32_t groups = (width + 5) / 6;

        for (uint32_t g = 0; g < groups; ++g)
        {
            uint32_t cb0 = (g * 6 + 0 < width) ? pixels[g * 6 + 0].cb : pixels[width - 1].cb;
            uint32_t y0  = (g * 6 + 0 < width) ? pixels[g * 6 + 0].y  : pixels[width - 1].y;
            uint32_t cr0 = (g * 6 + 0 < width) ? pixels[g * 6 + 0].cr : pixels[width - 1].cr;
            uint32_t y1  = (g * 6 + 1 < width) ? pixels[g * 6 + 1].y  : pixels[width - 1].y;
            uint32_t cb2 = (g * 6 + 2 < width) ? pixels[g * 6 + 2].cb : pixels[width - 1].cb;
            uint32_t y2  = (g * 6 + 2 < width) ? pixels[g * 6 + 2].y  : pixels[width - 1].y;
            uint32_t cr2 = (g * 6 + 3 < width) ? pixels[g * 6 + 3].cr : pixels[width - 1].cr;
            uint32_t y3  = (g * 6 + 3 < width) ? pixels[g * 6 + 3].y  : pixels[width - 1].y;
            uint32_t cb4 = (g * 6 + 4 < width) ? pixels[g * 6 + 4].cb : pixels[width - 1].cb;
            uint32_t y4  = (g * 6 + 4 < width) ? pixels[g * 6 + 4].y  : pixels[width - 1].y;
            uint32_t cr4 = (g * 6 + 5 < width) ? pixels[g * 6 + 5].cr : pixels[width - 1].cr;
            uint32_t y5  = (g * 6 + 5 < width) ? pixels[g * 6 + 5].y  : pixels[width - 1].y;

            uint32_t w0 = (cb0 & 0x3FF)
                        | ((y0 & 0x3FF) << 10)
                        | ((cr0 & 0x3FF) << 20);
            uint32_t w1 = (y1 & 0x3FF)
                        | ((cb2 & 0x3FF) << 10)
                        | ((y2 & 0x3FF) << 20);
            uint32_t w2 = (cr2 & 0x3FF)
                        | ((y3 & 0x3FF) << 10)
                        | ((cb4 & 0x3FF) << 20);
            uint32_t w3 = (y4 & 0x3FF)
                        | ((cr4 & 0x3FF) << 10)
                        | ((y5 & 0x3FF) << 20);

            uint8_t* p = out + g * 16;
            p[0]  = static_cast<uint8_t>(w0);
            p[1]  = static_cast<uint8_t>(w0 >> 8);
            p[2]  = static_cast<uint8_t>(w0 >> 16);
            p[3]  = static_cast<uint8_t>(w0 >> 24);
            p[4]  = static_cast<uint8_t>(w1);
            p[5]  = static_cast<uint8_t>(w1 >> 8);
            p[6]  = static_cast<uint8_t>(w1 >> 16);
            p[7]  = static_cast<uint8_t>(w1 >> 24);
            p[8]  = static_cast<uint8_t>(w2);
            p[9]  = static_cast<uint8_t>(w2 >> 8);
            p[10] = static_cast<uint8_t>(w2 >> 16);
            p[11] = static_cast<uint8_t>(w2 >> 24);
            p[12] = static_cast<uint8_t>(w3);
            p[13] = static_cast<uint8_t>(w3 >> 8);
            p[14] = static_cast<uint8_t>(w3 >> 16);
            p[15] = static_cast<uint8_t>(w3 >> 24);
        }
    }
}

int main()
{
    constexpr uint32_t width = 1920;
    constexpr uint32_t height = 1;
    constexpr uint32_t dstWidth = 512;
    constexpr uint32_t dstHeight = 1;

    uint32_t groups = (width + 5) / 6;
    uint32_t stride = ((groups * 16 + 127) / 128) * 128;

    printf("v210 conversion test\n");
    printf("width=%u stride=%u\n", width, stride);

    std::vector<YCbCr10> line(width);

    for (uint32_t x = 0; x < width; ++x)
    {
        size_t bar = (x * 8) / width;
        if (bar > 7) bar = 7;
        line[x] = rgbToYCbCr10(bars[bar].rgb75);
    }

    printf("\nReference YCbCr values:\n");
    for (size_t i = 0; i < 8; ++i)
    {
        YCbCr10 ref = rgbToYCbCr10(bars[i].rgb75);
        printf("  %-10s Y=%4u Cb=%4u Cr=%4u\n",
               bars[i].name, ref.y, ref.cb, ref.cr);
    }

    std::vector<uint8_t> v210Data(stride, 0);
    encodeV210Line(line.data(), width, v210Data.data());

    std::vector<uint32_t> rgbaBuf(dstWidth * dstHeight, 0);

    convertV210ToRgba(
        v210Data.data(),
        width,
        height,
        stride,
        rgbaBuf.data(),
        dstWidth,
        dstHeight);

    printf("\nDecode results:\n");
    int totalErrors = 0;

    for (size_t i = 0; i < 8; ++i)
    {
        uint32_t x = (i * dstWidth) / 8 + dstWidth / 16;
        if (x >= dstWidth) x = dstWidth - 1;

        uint32_t pixel = rgbaBuf[x];
        uint8_t r = static_cast<uint8_t>(pixel & 0xFF);
        uint8_t g = static_cast<uint8_t>((pixel >> 8) & 0xFF);
        uint8_t b = static_cast<uint8_t>((pixel >> 16) & 0xFF);

        RgbColor expected = bars[i].rgb75;

        int dr = static_cast<int>(r) - static_cast<int>(expected.r);
        int dg = static_cast<int>(g) - static_cast<int>(expected.g);
        int db = static_cast<int>(b) - static_cast<int>(expected.b);

        bool ok = (dr >= -3 && dr <= 3 && dg >= -3 && dg <= 3 && db >= -3 && db <= 3);

        if (!ok) totalErrors++;

        printf("  %-10s: got (%3u,%3u,%3u) expected (%3u,%3u,%3u) delta=(%+d,%+d,%+d) %s\n",
               bars[i].name, r, g, b,
               expected.r, expected.g, expected.b,
               dr, dg, db,
               ok ? "OK" : "FAIL");
    }

    printf("\n%d bar(s) out of tolerance\n", totalErrors);

    return totalErrors > 0 ? 1 : 0;
}