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

#define BUFFER_SIZE 128
#define FILE_HEADER_STRING "__info2_image_file_format__"

/* Hilfsfunktionen (module-local) */

static int checkFileHeader(FILE *fp)
{
    size_t headerLen = strlen(FILE_HEADER_STRING);
    if (headerLen >= BUFFER_SIZE) return 0;
    char buf[BUFFER_SIZE] = {0};
    if (fread(buf, sizeof(char), headerLen, fp) != headerLen) return 0;
    return (memcmp(buf, FILE_HEADER_STRING, headerLen) == 0);
}

static int readUShort(FILE *fp, unsigned short *out)
{
    if (!fp || !out) return 0;
    uint16_t tmp = 0;
    if (fread(&tmp, sizeof(uint16_t), 1, fp) != 1) return 0;
    *out = (unsigned short)tmp;
    return 1;
}

static int readPixels(FILE *fp, unsigned char *buffer, size_t count)
{
    if (count == 0) return 1;
    if (fread(buffer, sizeof(unsigned char), count, fp) != count) return 0;
    return 1;
}

static int readLabel(FILE *fp, unsigned char *outLabel)
{
    if (!fp || !outLabel) return 0;
    if (fread(outLabel, sizeof(unsigned char), 1, fp) != 1) return 0;
    return 1;
}

static void freeImagesInternal(GrayScaleImage *images, unsigned int count)
{
    if (!images) return;
    for (unsigned int i = 0; i < count; ++i) {
        free(images[i].buffer);
        images[i].buffer = NULL;
    }
    free(images);
}

/* Hauptfunktion: readImages im binären Format */
GrayScaleImageSeries *readImages(const char *path)
{
    if (!path) return NULL;

    FILE *fp = fopen(path, "rb");
    if (!fp) return NULL;

    if (!checkFileHeader(fp)) {
        fclose(fp);
        return NULL;
    }

    unsigned short count = 0, width = 0, height = 0;
    if (!readUShort(fp, &count) || !readUShort(fp, &width) || !readUShort(fp, &height)) {
        fclose(fp);
        return NULL;
    }

    GrayScaleImageSeries *series = (GrayScaleImageSeries *)malloc(sizeof(GrayScaleImageSeries));
    if (!series) { fclose(fp); return NULL; }

    series->count = (unsigned int)count;
    series->images = NULL;
    series->labels = NULL;

    if (count == 0) {
        /* leere Serie */
        fclose(fp);
        return series;
    }

    series->images = (GrayScaleImage *)calloc((size_t)count, sizeof(GrayScaleImage));
    series->labels = (unsigned char *)calloc((size_t)count, sizeof(unsigned char));
    if (!series->images || !series->labels) {
        free(series->images);
        free(series->labels);
        free(series);
        fclose(fp);
        return NULL;
    }

    size_t pixelCount = (size_t)width * (size_t)height;

    for (unsigned int i = 0; i < (unsigned int)count; ++i) {
        /* Pixeldaten zuerst (width*height bytes), dann Label (1 byte) */
        unsigned char *pixels = NULL;
        if (pixelCount > 0) {
            pixels = (unsigned char *)malloc(pixelCount * sizeof(unsigned char));
            if (!pixels) {
                /* cleanup */
                freeImagesInternal(series->images, i);
                free(series->labels);
                free(series);
                fclose(fp);
                return NULL;
            }

            if (!readPixels(fp, pixels, pixelCount)) {
                free(pixels);
                freeImagesInternal(series->images, i);
                free(series->labels);
                free(series);
                fclose(fp);
                return NULL;
            }
        }

        unsigned char label = 0;
        if (!readLabel(fp, &label)) {
            free(pixels);
            freeImagesInternal(series->images, i);
            free(series->labels);
            free(series);
            fclose(fp);
            return NULL;
        }

        series->images[i].buffer = pixels;
        series->images[i].width = width;
        series->images[i].height = height;
        series->images[i].label = label;
        series->labels[i] = label;
    }

    fclose(fp);
    return series;
}

/* Freigabe der Serie */
void clearSeries(GrayScaleImageSeries *series)
{
    if (!series) return;
    if (series->images) {
        for (unsigned int i = 0; i < series->count; ++i) {
            free(series->images[i].buffer);
            series->images[i].buffer = NULL;
        }
        free(series->images);
        series->images = NULL;
    }
    if (series->labels) {
        free(series->labels);
        series->labels = NULL;
    }
    free(series);
}