#include <unity.h>
#include <cstdint>
#include <cmath>
#include <cstdlib>
#include <cstring>
#include <climits>

// Palette definition (same as config.h)
struct Color { uint8_t r, g, b; };

static const Color PALETTE[6] = {
    {0, 0, 0},         // Black
    {255, 255, 255},   // White
    {200, 30, 30},     // Red
    {30, 160, 30},     // Green
    {30, 30, 200},     // Blue
    {220, 200, 30}     // Yellow
};

// Find nearest palette color (Euclidean distance in RGB)
uint8_t findNearestColor(int r, int g, int b) {
    uint8_t best = 0;
    int bestDist = INT32_MAX;
    for (int i = 0; i < 6; i++) {
        int dr = r - PALETTE[i].r;
        int dg = g - PALETTE[i].g;
        int db = b - PALETTE[i].b;
        int dist = dr*dr + dg*dg + db*db;
        if (dist < bestDist) {
            bestDist = dist;
            best = static_cast<uint8_t>(i);
        }
    }
    return best;
}

// Floyd-Steinberg dithering on a small test buffer
void ditherBuffer(uint8_t* rgb, int width, int height, uint8_t* output) {
    // Working buffer with int16_t to handle error diffusion overflow
    int16_t* work = (int16_t*)malloc(width * height * 3 * sizeof(int16_t));
    for (int i = 0; i < width * height * 3; i++) {
        work[i] = rgb[i];
    }

    for (int y = 0; y < height; y++) {
        for (int x = 0; x < width; x++) {
            int idx = (y * width + x) * 3;
            int r = work[idx];
            int g = work[idx + 1];
            int b = work[idx + 2];

            // Clamp
            r = r < 0 ? 0 : (r > 255 ? 255 : r);
            g = g < 0 ? 0 : (g > 255 ? 255 : g);
            b = b < 0 ? 0 : (b > 255 ? 255 : b);

            uint8_t nearest = findNearestColor(r, g, b);
            output[y * width + x] = nearest;

            // Error
            int errR = r - PALETTE[nearest].r;
            int errG = g - PALETTE[nearest].g;
            int errB = b - PALETTE[nearest].b;

            // Distribute error (Floyd-Steinberg weights: 7/16, 3/16, 5/16, 1/16)
            if (x + 1 < width) {
                int ni = (y * width + (x + 1)) * 3;
                work[ni]     += errR * 7 / 16;
                work[ni + 1] += errG * 7 / 16;
                work[ni + 2] += errB * 7 / 16;
            }
            if (y + 1 < height) {
                if (x > 0) {
                    int ni = ((y + 1) * width + (x - 1)) * 3;
                    work[ni]     += errR * 3 / 16;
                    work[ni + 1] += errG * 3 / 16;
                    work[ni + 2] += errB * 3 / 16;
                }
                {
                    int ni = ((y + 1) * width + x) * 3;
                    work[ni]     += errR * 5 / 16;
                    work[ni + 1] += errG * 5 / 16;
                    work[ni + 2] += errB * 5 / 16;
                }
                if (x + 1 < width) {
                    int ni = ((y + 1) * width + (x + 1)) * 3;
                    work[ni]     += errR * 1 / 16;
                    work[ni + 1] += errG * 1 / 16;
                    work[ni + 2] += errB * 1 / 16;
                }
            }
        }
    }
    free(work);
}

void test_nearest_color_black() {
    TEST_ASSERT_EQUAL(0, findNearestColor(0, 0, 0));
}

void test_nearest_color_white() {
    TEST_ASSERT_EQUAL(1, findNearestColor(255, 255, 255));
}

void test_nearest_color_red() {
    TEST_ASSERT_EQUAL(2, findNearestColor(180, 20, 20));
}

void test_nearest_color_green() {
    TEST_ASSERT_EQUAL(3, findNearestColor(20, 140, 20));
}

void test_nearest_color_blue() {
    TEST_ASSERT_EQUAL(4, findNearestColor(20, 20, 180));
}

void test_nearest_color_yellow() {
    TEST_ASSERT_EQUAL(5, findNearestColor(200, 180, 20));
}

void test_dither_solid_black() {
    const int W = 4, H = 4;
    uint8_t rgb[W * H * 3] = {0};  // All black
    uint8_t output[W * H];
    ditherBuffer(rgb, W, H, output);
    for (int i = 0; i < W * H; i++) {
        TEST_ASSERT_EQUAL(0, output[i]);  // All should be black
    }
}

void test_dither_solid_white() {
    const int W = 4, H = 4;
    uint8_t rgb[W * H * 3];
    memset(rgb, 255, sizeof(rgb));  // All white
    uint8_t output[W * H];
    ditherBuffer(rgb, W, H, output);
    for (int i = 0; i < W * H; i++) {
        TEST_ASSERT_EQUAL(1, output[i]);  // All should be white
    }
}

void test_dither_produces_valid_indices() {
    const int W = 8, H = 8;
    uint8_t rgb[W * H * 3];
    // Fill with mid-gray
    for (int i = 0; i < W * H * 3; i++) rgb[i] = 128;
    uint8_t output[W * H];
    ditherBuffer(rgb, W, H, output);
    for (int i = 0; i < W * H; i++) {
        TEST_ASSERT_TRUE(output[i] < 6);  // Valid palette index
    }
}