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

// Calibrated Spectra 6 palette (matches PALETTE_CALIBRATED in image_pipeline.cpp)
struct Color { uint8_t r, g, b; };

static const Color PALETTE[6] = {
    {0x1F, 0x22, 0x26},  // Black  -> dark gray
    {0xB9, 0xC7, 0xC9},  // White  -> light gray-blue
    {0x62, 0x20, 0x1E},  // Red    -> dark red/brown
    {0x35, 0x56, 0x3A},  // Green  -> dark forest green
    {0x23, 0x3F, 0x8E},  // Blue   -> dark navy
    {0xC1, 0xBB, 0x1E}   // Yellow -> olive/mustard
};

// Find nearest palette color (Rec. 709 luminance-weighted RGB distance)
uint8_t findNearestColor(int r, int g, int b) {
    uint8_t best = 0;
    int32_t 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;
        int32_t dist = 2126 * dr * dr + 7152 * dg * dg + 722 * 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
    }
}

// Include the actual blue noise texture for testing
#include "../../src/blue_noise.h"

static constexpr float TEST_BLUE_NOISE_STRENGTH = 32.0f;

// Floyd-Steinberg dither with blue noise threshold modulation
void ditherBufferWithBlueNoise(uint8_t* rgb, int width, int height, uint8_t* output) {
    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];

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

            // Apply blue noise perturbation
            uint8_t noise = pgm_read_byte(&BLUE_NOISE_64[(y & 63) * 64 + (x & 63)]);
            float nf = ((float)noise - 128.0f) * (TEST_BLUE_NOISE_STRENGTH / 128.0f);
            int rn = r + (int)nf;
            int gn = g + (int)nf;
            int bn = b + (int)nf;
            rn = rn < 0 ? 0 : (rn > 255 ? 255 : rn);
            gn = gn < 0 ? 0 : (gn > 255 ? 255 : gn);
            bn = bn < 0 ? 0 : (bn > 255 ? 255 : bn);

            uint8_t nearest = findNearestColor(rn, gn, bn);
            output[y * width + x] = nearest;

            // Error uses original r,g,b (energy conservation)
            int errR = r - PALETTE[nearest].r;
            int errG = g - PALETTE[nearest].g;
            int errB = b - PALETTE[nearest].b;

            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_blue_noise_breaks_pattern_regularity() {
    // Use a 16x16 dark gray patch — just above the black palette value.
    // Standard F-S produces highly regular dot patterns in this region.
    const int W = 16, H = 16;
    uint8_t rgb[W * H * 3];
    memset(rgb, 45, sizeof(rgb));  // Dark gray (45,45,45) — triggers error accumulation

    uint8_t output_standard[W * H];
    uint8_t output_blue_noise[W * H];

    // Copy rgb since ditherBuffer modifies working buffer
    uint8_t rgb_copy[W * H * 3];
    memcpy(rgb_copy, rgb, sizeof(rgb));

    ditherBuffer(rgb, W, H, output_standard);
    ditherBufferWithBlueNoise(rgb_copy, W, H, output_blue_noise);

    // Both should produce valid palette indices
    for (int i = 0; i < W * H; i++) {
        TEST_ASSERT_TRUE(output_standard[i] < 6);
        TEST_ASSERT_TRUE(output_blue_noise[i] < 6);
    }

    // Count unique rows in each output to measure pattern regularity.
    // Standard F-S in uniform areas tends to produce repeating row patterns.
    // Blue noise should produce more unique rows (less periodic).
    int unique_standard = 0;
    int unique_blue_noise = 0;

    for (int y = 0; y < H; y++) {
        bool is_duplicate = false;
        for (int prev = 0; prev < y; prev++) {
            if (memcmp(&output_standard[y * W], &output_standard[prev * W], W) == 0) {
                is_duplicate = true;
                break;
            }
        }
        if (!is_duplicate) unique_standard++;
    }

    for (int y = 0; y < H; y++) {
        bool is_duplicate = false;
        for (int prev = 0; prev < y; prev++) {
            if (memcmp(&output_blue_noise[y * W], &output_blue_noise[prev * W], W) == 0) {
                is_duplicate = true;
                break;
            }
        }
        if (!is_duplicate) unique_blue_noise++;
    }

    // Blue noise version should have at least as many unique rows as standard.
    // In practice it should have more (the whole point of the feature).
    TEST_ASSERT_GREATER_OR_EQUAL(unique_standard, unique_blue_noise);
}

void test_blue_noise_dither_valid_indices() {
    const int W = 16, H = 16;
    uint8_t rgb[W * H * 3];
    // Test across different gray levels
    for (int level = 0; level < 256; level += 32) {
        memset(rgb, level, sizeof(rgb));
        uint8_t output[W * H];
        ditherBufferWithBlueNoise(rgb, W, H, output);
        for (int i = 0; i < W * H; i++) {
            TEST_ASSERT_TRUE(output[i] < 6);
        }
    }
}