#include "test_common.h"
#include <stdio.h>
#include <stdlib.h>

int test_failures = 0;
int test_total = 0;

static int test_null_image_handling(void) {
    uint8_t *result = downsample_image(NULL, 100, 100, 4, NULL, NULL, 0.5f);
    TEST_ASSERT(result == NULL, "Null image should return NULL");
    return TEST_PASSED;
}

static int test_zero_dimensions(void) {
    int dw, dh;
    uint8_t *result = downsample_image(NULL, 0, 0, 4, &dw, &dh, 0.5f);
    TEST_ASSERT(result == NULL, "Zero dimensions should return NULL");
    return TEST_PASSED;
}

static int test_scale_one_preserves_size(void) {
    uint8_t *src = create_uniform_image(100, 100, 128, 64, 255);
    TEST_ASSERT_PTR_NOT_NULL(src, "Source image allocation");

    int dw, dh;
    uint8_t *dst = downsample_image(src, 100, 100, 4, &dw, &dh, 1.0f);
    TEST_ASSERT_PTR_NOT_NULL(dst, "Scale 1.0 result");

    TEST_ASSERT_EQ(dw, 100, "Width preserved at scale 1.0");
    TEST_ASSERT_EQ(dh, 100, "Height preserved at scale 1.0");

    int match = compare_images(src, dst, 100, 100, 4, 0.01f);
    TEST_ASSERT(match == 1, "Image data preserved at scale 1.0");

    free(src);
    free(dst);
    return TEST_PASSED;
}

static int test_downsample_half_size(void) {
    uint8_t *src = create_gradient_image(100, 100);
    TEST_ASSERT_PTR_NOT_NULL(src, "Gradient source");

    int dw, dh;
    uint8_t *dst = downsample_image(src, 100, 100, 4, &dw, &dh, 0.5f);
    TEST_ASSERT_PTR_NOT_NULL(dst, "Downsample result");

    TEST_ASSERT_EQ(dw, 50, "Width halved at 0.5 scale");
    TEST_ASSERT_EQ(dh, 50, "Height halved at 0.5 scale");

    free(src);
    free(dst);
    return TEST_PASSED;
}

static int test_minimum_one_pixel(void) {
    uint8_t *src = create_uniform_image(1, 1, 100, 100, 100);
    TEST_ASSERT_PTR_NOT_NULL(src, "Single pixel source");

    int dw, dh;
    uint8_t *dst = downsample_image(src, 1, 1, 4, &dw, &dh, 0.01f);
    TEST_ASSERT_PTR_NOT_NULL(dst, "Minimum size result");

    TEST_ASSERT_EQ(dw, 1, "Width minimum 1 pixel");
    TEST_ASSERT_EQ(dh, 1, "Height minimum 1 pixel");

    free(src);
    free(dst);
    return TEST_PASSED;
}

static int test_large_image_handling(void) {
    uint8_t *src = create_noise_image(4096, 4096, 42);
    TEST_ASSERT_PTR_NOT_NULL(src, "Large image allocation");

    int dw, dh;
    uint8_t *dst = downsample_image(src, 4096, 4096, 4, &dw, &dh, 0.25f);
    TEST_ASSERT_PTR_NOT_NULL(dst, "Large image downsample");

    TEST_ASSERT_EQ(dw, 1024, "Correct width at 0.25 scale");
    TEST_ASSERT_EQ(dh, 1024, "Correct height at 0.25 scale");

    free(src);
    free(dst);
    return TEST_PASSED;
}

static int test_alpha_channel_preserved(void) {
    uint8_t *src = create_uniform_image(50, 50, 255, 0, 0);
    TEST_ASSERT_PTR_NOT_NULL(src, "Red source");

    int dw, dh;
    uint8_t *dst = downsample_image(src, 50, 50, 4, &dw, &dh, 0.5f);
    TEST_ASSERT_PTR_NOT_NULL(dst, "Downsample with alpha");

    TEST_ASSERT_EQ(dw, 25, "Width halved");
    TEST_ASSERT_EQ(dh, 25, "Height halved");

    TEST_ASSERT(dst[0 * 4 + 3] == 255, "Alpha channel preserved");

    free(src);
    free(dst);
    return TEST_PASSED;
}

static int test_uniform_color_accuracy(void) {
    uint8_t *src = create_uniform_image(100, 100, 200, 100, 50);
    TEST_ASSERT_PTR_NOT_NULL(src, "Uniform color source");

    int dw, dh;
    uint8_t *dst = downsample_image(src, 100, 100, 4, &dw, &dh, 0.25f);
    TEST_ASSERT_PTR_NOT_NULL(dst, "Uniform color downsample");

    int correct = 1;
    for (int i = 0; i < dw * dh; i++) {
        if (dst[i * 4 + 0] != 200 || dst[i * 4 + 1] != 100 || dst[i * 4 + 2] != 50) {
            correct = 0;
            break;
        }
    }
    TEST_ASSERT(correct == 1, "Uniform color preserved in downsampling");

    free(src);
    free(dst);
    return TEST_PASSED;
}

static int test_gradient_smoothness(void) {
    uint8_t *src = create_gradient_image(100, 100);
    TEST_ASSERT_PTR_NOT_NULL(src, "Gradient source");

    int dw, dh;
    uint8_t *dst = downsample_image(src, 100, 100, 4, &dw, &dh, 0.5f);
    TEST_ASSERT_PTR_NOT_NULL(dst, "Gradient downsample");

    float psnr = compute_psnr(src, dst, dw, dh, 4);
    TEST_ASSERT_FTZ(psnr, 9.0f, 5.0f, "Gradient PSNR within expected range");

    free(src);
    free(dst);
    return TEST_PASSED;
}

static int test_checkerboard_no_aliasing(void) {
    uint8_t *src = create_checkerboard(100, 100, 10);
    TEST_ASSERT_PTR_NOT_NULL(src, "Checkerboard source");

    int dw, dh;
    uint8_t *dst = downsample_image(src, 100, 100, 4, &dw, &dh, 0.5f);
    TEST_ASSERT_PTR_NOT_NULL(dst, "Checkerboard downsample");

    TEST_ASSERT_EQ(dw, 50, "Width halved");
    TEST_ASSERT_EQ(dh, 50, "Height halved");

    free(src);
    free(dst);
    return TEST_PASSED;
}

static TestCase tests[] = {
    {"null_image_handling", test_null_image_handling},
    {"zero_dimensions", test_zero_dimensions},
    {"scale_one_preserves_size", test_scale_one_preserves_size},
    {"downsample_half_size", test_downsample_half_size},
    {"minimum_one_pixel", test_minimum_one_pixel},
    {"large_image_handling", test_large_image_handling},
    {"alpha_channel_preserved", test_alpha_channel_preserved},
    {"uniform_color_accuracy", test_uniform_color_accuracy},
    {"gradient_smoothness", test_gradient_smoothness},
    {"checkerboard_no_aliasing", test_checkerboard_no_aliasing},
};

int main(int argc, char **argv) {
    int profile = 0;
    for (int i = 1; i < argc; i++) {
        if (strcmp(argv[i], "--profile") == 0) {
            profile = 1;
        }
    }

    if (profile) {
        const char *scenarios[] = {"No_Downsampling", "1080p_Target", "1440p_Target", "4K_Target"};
        float scales[] = {1.0f, 0.5f, 0.42f, 0.25f};

        for (int s = 0; s < 4; s++) {
            char filename[256];
            snprintf(filename, sizeof(filename), "/tmp/chroma_memory_%s.csv", scenarios[s]);

            FILE *f = fopen(filename, "w");
            if (!f) {
                fprintf(stderr, "Failed to create %s\n", filename);
                continue;
            }

            fprintf(f, "Resolution,OriginalSizeMB,DownsampledSizeMB,MemorySavingsPercent\n");

            const char *res_names[] = {"1080p", "1440p", "4K", "5K", "6K", "8K"};
            int widths[] = {1920, 2560, 3840, 5120, 6016, 7680};
            int heights[] = {1080, 1440, 2160, 2880, 3200, 4320};

            for (int r = 0; r < 6; r++) {
                int w = widths[r];
                int h = heights[r];
                size_t original_bytes = (size_t)w * (size_t)h * 4u;
                double original_mb = (double)original_bytes / (1024.0 * 1024.0);

                int dw, dh;
                uint8_t *src = create_noise_image(w, h, 42);
                uint8_t *dst = downsample_image(src, w, h, 4, &dw, &dh, scales[s]);
                size_t downsampled_bytes = (size_t)dw * (size_t)dh * 4u;
                double downsampled_mb = (double)downsampled_bytes / (1024.0 * 1024.0);

                double savings = 0.0;
                if (original_mb > 0) {
                    savings = ((original_mb - downsampled_mb) / original_mb) * 100.0;
                }

                fprintf(f, "%s,%.2f,%.2f,%.1f\n", res_names[r], original_mb, downsampled_mb, savings);

                free(src);
                free(dst);
            }

            fclose(f);
            printf("Generated: %s\n", filename);
        }
        return 0;
    }

    (void)argc;
    (void)argv;

    printf("Chroma Unit Tests\n");
    printf("=================\n\n");

    test_failures = 0;
    test_total = 0;

    for (int i = 0; i < (int)(sizeof(tests) / sizeof(tests[0])); i++) {
        int result = tests[i].fn();
        test_total++;
        if (result == TEST_PASSED) {
            printf("  [PASS] %s\n", tests[i].name);
        } else {
            printf("  [FAIL] %s\n", tests[i].name);
        }
    }

    printf("\n-----------------\n");
    printf("Results: %d/%d passed\n", test_total - test_failures, test_total);

    return test_failures > 0 ? 1 : 0;
}