#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include "unity.h"
#include "imageInput.h"


static void writeHeader(FILE *file, unsigned short count, unsigned short width, unsigned short height)
{
    const char *tag = "__info2_image_file_format__";
    fwrite(tag, 1, strlen(tag), file);
    fwrite(&count, sizeof(count), 1, file);
    fwrite(&width, sizeof(width), 1, file);
    fwrite(&height, sizeof(height), 1, file);
}

static void prepareImageFile(const char *path, unsigned short width, unsigned short height,
                             unsigned short count, unsigned char label)
{
    FILE *file = fopen(path, "wb");
    if (file == NULL) return;

    writeHeader(file, count, width, height);
    size_t pixelCount = (size_t)width * height;
    GrayScalePixelType *pixels = (GrayScalePixelType *)calloc(pixelCount, sizeof(*pixels));
    if (pixels != NULL)
    {
        for (unsigned int i = 0; i < count; ++i)
        {
            fwrite(pixels, sizeof(*pixels), pixelCount, file);
            fwrite(&label, sizeof(label), 1, file);
        }
        free(pixels);
    }
    fclose(file);
}

void test_readImagesReturnsCorrectNumberOfImages(void)
{
    const char *path = "testFile.info2";
    const unsigned short expectedCount = 2;
    prepareImageFile(path, 8, 8, expectedCount, 1);
    GrayScaleImageSeries *series = readImages(path);
    TEST_ASSERT_NOT_NULL(series);
    TEST_ASSERT_EQUAL_UINT16(expectedCount, series->count);
    clearSeries(series);
    remove(path);
}

void test_readImagesReturnsCorrectImageDimensions(void)
{
    const char *path = "testFile.info2";
    prepareImageFile(path, 10, 9, 2, 1);
    GrayScaleImageSeries *series = readImages(path);
    TEST_ASSERT_NOT_NULL(series);
    TEST_ASSERT_EQUAL_UINT16(10, series->images[0].width);
    TEST_ASSERT_EQUAL_UINT16(10, series->images[1].width);
    TEST_ASSERT_EQUAL_UINT16(9, series->images[0].height);
    TEST_ASSERT_EQUAL_UINT16(9, series->images[1].height);
    clearSeries(series);
    remove(path);
}

void test_readImagesReturnsCorrectLabels(void)
{
    const char *path = "testFile.info2";
    prepareImageFile(path, 8, 8, 2, 15);
    GrayScaleImageSeries *series = readImages(path);
    TEST_ASSERT_NOT_NULL(series);
    TEST_ASSERT_EQUAL_UINT8(15, series->labels[0]);
    TEST_ASSERT_EQUAL_UINT8(15, series->labels[1]);
    clearSeries(series);
    remove(path);
}

void test_readImagesReturnsPixelData(void)
{
    const char *path = "testFile.info2";
    const unsigned char pixels[] = { 0, 17, 128, 255 };
    const unsigned char label = 7;
    FILE *file = fopen(path, "wb");
    TEST_ASSERT_NOT_NULL(file);
    writeHeader(file, 1, 2, 2);
    fwrite(pixels, sizeof(pixels), 1, file);
    fwrite(&label, sizeof(label), 1, file);
    fclose(file);

    GrayScaleImageSeries *series = readImages(path);
    TEST_ASSERT_NOT_NULL(series);
    TEST_ASSERT_EQUAL_UINT8_ARRAY(pixels, series->images[0].buffer, sizeof(pixels));
    clearSeries(series);
    remove(path);
}

void test_readImagesReturnsNullOnMissingPath(void)
{
    const char *path = "missingTestFile.info2";
    remove(path);
    TEST_ASSERT_NULL(readImages(path));
}

void test_readImagesFailsOnWrongFileTag(void)
{
    const char *path = "testFile.info2";
    FILE *file = fopen(path, "wb");
    TEST_ASSERT_NOT_NULL(file);
    fputs("wrong tag", file);
    fclose(file);
    TEST_ASSERT_NULL(readImages(path));
    remove(path);
}

void test_readImagesFailsOnTruncatedImage(void)
{
    const char *path = "testFile.info2";
    const unsigned char partialPixels[] = { 1, 2 };
    FILE *file = fopen(path, "wb");
    TEST_ASSERT_NOT_NULL(file);
    writeHeader(file, 1, 2, 2);
    fwrite(partialPixels, sizeof(partialPixels), 1, file);
    fclose(file);
    TEST_ASSERT_NULL(readImages(path));
    remove(path);
}

void setUp(void) {}
void tearDown(void) {}

int main(void)
{
    UNITY_BEGIN();
    RUN_TEST(test_readImagesReturnsCorrectNumberOfImages);
    RUN_TEST(test_readImagesReturnsCorrectImageDimensions);
    RUN_TEST(test_readImagesReturnsCorrectLabels);
    RUN_TEST(test_readImagesReturnsPixelData);
    RUN_TEST(test_readImagesReturnsNullOnMissingPath);
    RUN_TEST(test_readImagesFailsOnWrongFileTag);
    RUN_TEST(test_readImagesFailsOnTruncatedImage);
    return UNITY_END();
}