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

#define BUFFER_SIZE 100
#define FILE_HEADER_STRING "__info2_image_file_format__"

typedef enum
{
    IMG_SUCCESS = 0,
    IMG_ERR_INVALID_HEADER, // Header did not match
    IMG_ERR_READ,           // Failed to read from file, maybe it's too short
    IMG_ERR,                // General Error
} ImageError;

static ImageError checkHeader(FILE *file);
static ImageError readPictureParams(unsigned short *number, unsigned short *width, unsigned short *height, FILE *file);
static ImageError readImage(size_t numPixels, GrayScalePixelType *pixelBuffer, unsigned char *label, FILE *file);
static ImageError parseImageFile(FILE *file, GrayScaleImageSeries *series);

GrayScaleImageSeries *readImages(const char *path)
{
    // It's very important to call calloc here because otherwise clearSeries() might try to free random memory
    GrayScaleImageSeries *series = calloc(1, sizeof(GrayScaleImageSeries));
    if (series == NULL)
    {
        return NULL;
    }

    FILE *file = fopen(path, "rb");
    // If fopen() failed
    if (file == NULL)
    {
        clearSeries(series);
        series = NULL;
        return NULL;
    }

    // Try to parse the whole file. If anything fails memory is freed
    if (parseImageFile(file, series))
    {
        clearSeries(series);
        series = NULL; // Return NULL after the file is properly closed 
    }

    fclose(file);
    return series;
}

static ImageError parseImageFile(FILE *file, GrayScaleImageSeries *series)
{
    if (checkHeader(file))
    {
        return IMG_ERR; // header check failed
    }

    unsigned short imgCNT, width, height;
    if (readPictureParams(&imgCNT, &width, &height, file))
    {
        return IMG_ERR; // read failed
    }
    size_t pixels = width * height;

    // The images contain more pointers, definitely use calloc here
    series->images = calloc(imgCNT, sizeof(GrayScaleImage));
    if (series->images == NULL)
    {
        return IMG_ERR;
    }

    series->labels = malloc(imgCNT * sizeof(unsigned char));
    if (series->labels == NULL)
    {
        return IMG_ERR;
    }

    series->count = imgCNT;
    // Read every image and it's label
    for (size_t imageIdx = 0; imageIdx < imgCNT; imageIdx++)
    {
        GrayScaleImage *curImage = &series->images[imageIdx];

        curImage->buffer = malloc(sizeof(GrayScalePixelType) * pixels);
        if (curImage->buffer == NULL)
        {
            return IMG_ERR;
        }

        curImage->width = width;
        curImage->height = height;

        if (readImage(pixels, curImage->buffer, &series->labels[imageIdx], file))
        {
            return IMG_ERR;
        }
    }
    return IMG_SUCCESS;
}

static ImageError checkHeader(FILE *file)
{
    size_t len = strlen(FILE_HEADER_STRING);
    char headerBuf[len + 1];
    size_t charsRead = fread(headerBuf, 1, len, file);
    // Check if the file is to short
    if (charsRead < len)
    {
        return IMG_ERR_READ;
    }
    headerBuf[len] = '\0'; // Terminate string
    return strcmp(headerBuf, FILE_HEADER_STRING) == 0 ? IMG_SUCCESS : IMG_ERR_INVALID_HEADER;
}

static ImageError readPictureParams(unsigned short *number, unsigned short *width, unsigned short *height, FILE *file)
{
    // number of images
    if (1 != fread(number, sizeof(unsigned short), 1, file))
    {
        return IMG_ERR_READ;
    }

    // image width
    if (1 != fread(width, sizeof(unsigned short), 1, file))
    {
        return IMG_ERR_READ;
    }

    // height
    if (1 != fread(height, sizeof(unsigned short), 1, file))
    {
        return IMG_ERR_READ;
    }

    return IMG_SUCCESS;
}

// Read 1 image and it's label
static ImageError readImage(size_t numPixels, GrayScalePixelType *pixelBuffer, unsigned char *label, FILE *file)
{
    if (numPixels > fread(pixelBuffer, sizeof(GrayScalePixelType), numPixels, file))
    {
        return IMG_ERR_READ;
    }

    if (1 != fread(label, sizeof(unsigned char), 1, file))
    {
        return IMG_ERR_READ;
    }

    return IMG_SUCCESS;
}

// Frees memory for each image buffer, image, label and finally series
void clearSeries(GrayScaleImageSeries *series)
{
    if (series)
    {

        int seriesLen = series->count;
        for (size_t imageIdx = 0; imageIdx < seriesLen; imageIdx++)
        {
            free(series->images[imageIdx].buffer);
        }

        free(series->images);
        free(series->labels);
        free(series);
    }
}