#include "image.h"

#include <SDL3_image/SDL_image.h>

#include <limits.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

typedef bool (*ImageProbe)(SDL_IOStream *stream);

static bool probe_file(const char *path, ImageProbe probe) {
    SDL_IOStream *stream = SDL_IOFromFile(path, "rb");
    bool result;
    if (stream == NULL) {
        return false;
    }
    result = probe(stream);
    SDL_CloseIO(stream);
    return result;
}

static bool file_has_magic(const char *path, const char *magic, size_t length) {
    uint8_t bytes[16];
    FILE *file;
    bool matches;
    if (length > sizeof(bytes)) {
        return false;
    }
    file = fopen(path, "rb");
    if (file == NULL) {
        return false;
    }
    matches = fread(bytes, 1, length, file) == length &&
              memcmp(bytes, magic, length) == 0;
    fclose(file);
    return matches;
}

static const char *animation_type(const char *path) {
    if (probe_file(path, IMG_isGIF)) return "GIF";
    if (probe_file(path, IMG_isWEBP)) return "WEBP";
    if (probe_file(path, IMG_isAVIF)) return "AVIF";
    if (probe_file(path, IMG_isPNG)) return "PNG";
    if (probe_file(path, IMG_isANI)) return "ANI";
    return NULL;
}

static bool add_surface_blocks(SDL_Surface *surface, uint32_t block_size,
                               uint64_t duration, CandidateSet *set) {
    uint32_t top;
    const double frame_weight = duration == 0 ? 1.0 : (double)duration;
    bool success = true;

    if (surface->w <= 0 || surface->h <= 0 || surface->pitch <= 0) {
        SDL_SetError("invalid surface dimensions or pitch");
        return false;
    }
    if (SDL_MUSTLOCK(surface) && !SDL_LockSurface(surface)) {
        return false;
    }
    for (top = 0; top < (uint32_t)surface->h && success; top += block_size) {
        uint32_t left;
        const uint32_t bottom = block_size < (uint32_t)surface->h - top
                                    ? top + block_size : (uint32_t)surface->h;
        for (left = 0; left < (uint32_t)surface->w && success;
             left += block_size) {
            const uint32_t right = block_size < (uint32_t)surface->w - left
                                       ? left + block_size : (uint32_t)surface->w;
            ColourAccumulator accumulator = {0.0, 0.0, 0.0, 0.0};
            RGB average;
            uint32_t y;
            for (y = top; y < bottom; ++y) {
                const uint8_t *pixel = (const uint8_t *)surface->pixels +
                                       (size_t)y * (size_t)surface->pitch +
                                       (size_t)left * 4U;
                uint32_t x;
                for (x = left; x < right; ++x) {
                    colour_accumulator_add_u8(&accumulator, pixel[0], pixel[1],
                                              pixel[2], pixel[3]);
                    pixel += 4;
                }
            }
            if (colour_accumulator_average(&accumulator, &average)) {
                success = candidates_add(set, average,
                                         accumulator.alpha * frame_weight);
            }
        }
    }
    if (SDL_MUSTLOCK(surface)) {
        SDL_UnlockSurface(surface);
    }
    return success;
}

bool image_process_surface(SDL_Surface *surface, uint32_t block_size,
                           uint64_t duration, CandidateSet *set) {
    SDL_Surface *normalised = surface;
    bool success;
    if (surface == NULL || block_size <= 0 || set == NULL) {
        SDL_SetError("invalid image processing arguments");
        return false;
    }
    if (surface->format != SDL_PIXELFORMAT_RGBA32) {
        normalised = SDL_ConvertSurface(surface, SDL_PIXELFORMAT_RGBA32);
        if (normalised == NULL) {
            return false;
        }
    }
    success = add_surface_blocks(normalised, block_size, duration, set);
    if (normalised != surface) {
        SDL_DestroySurface(normalised);
    }
    if (!success) {
        SDL_SetError("out of memory while collecting colours");
    }
    return success;
}

static uint32_t read_be32(const uint8_t *bytes) {
    return ((uint32_t)bytes[0] << 24) | ((uint32_t)bytes[1] << 16) |
           ((uint32_t)bytes[2] << 8) | (uint32_t)bytes[3];
}

static uint16_t read_be16(const uint8_t *bytes) {
    return (uint16_t)(((uint16_t)bytes[0] << 8) | (uint16_t)bytes[1]);
}

static bool process_farbfeld(const char *path, uint32_t block_size,
                             CandidateSet *set) {
    uint8_t header[16];
    uint8_t *band = NULL;
    Arena arena;
    FILE *file = fopen(path, "rb");
    uint32_t width;
    uint32_t height;
    uint32_t band_rows;
    uint32_t top;
    bool success = false;

    arena_init(&arena);
    if (file == NULL) {
        SDL_SetError("could not open farbfeld image");
        return false;
    }
    if (fread(header, 1, sizeof(header), file) != sizeof(header) ||
        memcmp(header, "farbfeld", 8) != 0) {
        SDL_SetError("invalid farbfeld header");
        goto done;
    }
    width = read_be32(header + 8);
    height = read_be32(header + 12);
    if (width == 0 || height == 0 || width > INT_MAX || height > INT_MAX) {
        SDL_SetError("invalid or excessively large farbfeld dimensions");
        goto done;
    }
    band_rows = block_size < height ? block_size : height;
    if ((size_t)width > SIZE_MAX / (size_t)band_rows) {
        SDL_SetError("farbfeld scanline band is too large");
        goto done;
    }
    band = arena_allocate(&arena, (size_t)width * (size_t)band_rows, 8U);
    if (band == NULL) {
        SDL_SetError("out of memory decoding farbfeld image");
        goto done;
    }
    for (top = 0; top < height; top += (uint32_t)block_size) {
        const uint32_t rows = block_size < height - top
                                  ? block_size : height - top;
        const size_t bytes = (size_t)width * (size_t)rows * 8U;
        uint32_t left;
        if (fread(band, 1, bytes, file) != bytes) {
            SDL_SetError("truncated farbfeld pixel data");
            goto done;
        }
        for (left = 0; left < width; left += block_size) {
            const uint32_t columns = block_size < width - left
                                         ? block_size : width - left;
            ColourAccumulator accumulator = {0.0, 0.0, 0.0, 0.0};
            RGB average;
            uint32_t row;
            for (row = 0; row < rows; ++row) {
                const uint8_t *pixel = band +
                    ((size_t)row * (size_t)width + left) * 8U;
                uint32_t column;
                for (column = 0; column < columns; ++column) {
                    colour_accumulator_add_u16(
                        &accumulator, read_be16(pixel), read_be16(pixel + 2),
                        read_be16(pixel + 4), read_be16(pixel + 6));
                    pixel += 8;
                }
            }
            if (colour_accumulator_average(&accumulator, &average) &&
                !candidates_add(set, average, accumulator.alpha)) {
                SDL_SetError("out of memory while collecting colours");
                goto done;
            }
        }
    }
    if (fgetc(file) != EOF) {
        /* Extra bytes are harmless and accepted by the farbfeld convention. */
    }
    success = true;
done:
    arena_destroy(&arena);
    fclose(file);
    return success;
}

static bool process_animation(const char *path, const char *type,
                              uint32_t block_size, CandidateSet *set,
                              bool *was_animation) {
    SDL_IOStream *stream = SDL_IOFromFile(path, "rb");
    IMG_AnimationDecoder *decoder;
    bool success = true;
    size_t frames = 0;

    *was_animation = false;
    if (stream == NULL) {
        return false;
    }
    SDL_ClearError();
    decoder = IMG_CreateAnimationDecoder_IO(stream, true, type);
    if (decoder == NULL) {
        SDL_ClearError();
        return true;
    }
    for (;;) {
        SDL_Surface *frame = NULL;
        Uint64 duration = 0;
        if (!IMG_GetAnimationDecoderFrame(decoder, &frame, &duration)) {
            const IMG_AnimationDecoderStatus status =
                IMG_GetAnimationDecoderStatus(decoder);
            if (status == IMG_DECODER_STATUS_FAILED ||
                status == IMG_DECODER_STATUS_INVALID) {
                success = false;
            }
            break;
        }
        ++frames;
        if (!image_process_surface(frame, block_size, duration, set)) {
            success = false;
            SDL_DestroySurface(frame);
            break;
        }
        SDL_DestroySurface(frame);
    }
    if (!IMG_CloseAnimationDecoder(decoder) && success) {
        success = false;
    }
    *was_animation = frames > 0;
    return success;
}

static bool process_static(const char *path, uint32_t block_size,
                           CandidateSet *set) {
    SDL_IOStream *stream = SDL_IOFromFile(path, "rb");
    SDL_Surface *surface;
    const char *extension;
    bool success;
    if (stream == NULL) {
        return false;
    }
    surface = IMG_Load_IO(stream, true);
    if (surface == NULL) {
        /* TGA has no reliable magic signature, so SDL_image cannot include it
           in normal content probing. Use the suffix only for this fallback. */
        extension = strrchr(path, '.');
        if (extension == NULL || SDL_strcasecmp(extension, ".tga") != 0) {
            return false;
        }
        stream = SDL_IOFromFile(path, "rb");
        if (stream == NULL) {
            return false;
        }
        surface = IMG_LoadTyped_IO(stream, true, "TGA");
        if (surface == NULL) {
            return false;
        }
    }
    success = image_process_surface(surface, block_size, 1, set);
    SDL_DestroySurface(surface);
    return success;
}

bool image_process_file(const char *path, uint32_t block_size,
                        CandidateSet *set) {
    const char *type;
    bool was_animation;
    if (path == NULL || block_size == 0 || set == NULL) {
        SDL_SetError("invalid image processing arguments");
        return false;
    }
    if (file_has_magic(path, "farbfeld", 8)) {
        return process_farbfeld(path, block_size, set);
    }
    if (probe_file(path, IMG_isSVG)) {
        SDL_SetError("SVG is intentionally unsupported");
        return false;
    }
    type = animation_type(path);
    if (type != NULL) {
        if (!process_animation(path, type, block_size, set, &was_animation)) {
            return false;
        }
        if (was_animation) {
            return true;
        }
    }
    return process_static(path, block_size, set);
}