#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <libheif/heif.h>

heif_image* create_tile(int w, int h)
{
    heif_image* img;
    heif_error err = heif_image_create(w, h, heif_colorspace_YCbCr, heif_chroma_420, &img);
    if (err.code != heif_error_Ok) {
        fprintf(stderr, "heif_image_create failed: %s\n", err.message);
        exit(1);
    }

    int cw = w/2 + (w & 1);
    int ch = h/2 + (h & 1);

    err = heif_image_add_plane(img, heif_channel_Y, w, h, 8);
    if (err.code != heif_error_Ok) { fprintf(stderr, "add Y failed: %s\n", err.message); exit(1); }
    err = heif_image_add_plane(img, heif_channel_Cb, cw, ch, 8);
    if (err.code != heif_error_Ok) { fprintf(stderr, "add Cb failed: %s\n", err.message); exit(1); }
    err = heif_image_add_plane(img, heif_channel_Cr, cw, ch, 8);
    if (err.code != heif_error_Ok) { fprintf(stderr, "add Cr failed: %s\n", err.message); exit(1); }

    int y_stride, cb_stride, cr_stride;
    uint8_t* Y = heif_image_get_plane(img, heif_channel_Y, &y_stride);
    uint8_t* Cb = heif_image_get_plane(img, heif_channel_Cb, &cb_stride);
    uint8_t* Cr = heif_image_get_plane(img, heif_channel_Cr, &cr_stride);

    memset(Y, 0x80, y_stride * h);
    memset(Cb, 0x41, cb_stride * ch);
    memset(Cr, 0x42, cr_stride * ch);

    return img;
}

int main(int argc, char** argv)
{
    if (argc < 2) {
        fprintf(stderr, "Usage: %s <output.heif> [decode]\n", argv[0]);
        return 1;
    }

    const char* out_file = argv[1];
    int decode_mode = (argc >= 3 && strcmp(argv[2], "decode") == 0);

    heif_context* ctx = heif_context_alloc();

    heif_image* tiles[4];
    for (int i = 0; i < 4; i++) {
        tiles[i] = create_tile(64, 65);
    }

    heif_encoder* encoder = NULL;
    heif_error err = heif_context_get_encoder_for_format(ctx, heif_compression_JPEG, &encoder);
    if (err.code != heif_error_Ok) {
        fprintf(stderr, "get_encoder failed: %s\n", err.message);
        return 1;
    }

    heif_image_handle* grid_handle = NULL;
    heif_encoding_options* opts = heif_encoding_options_alloc();
    err = heif_context_add_grid_image(ctx, 64, 260, 1, 4, opts, &grid_handle);
    heif_encoding_options_free(opts);
    if (err.code != heif_error_Ok) {
        fprintf(stderr, "add_grid_image failed: %s\n", err.message);
        return 1;
    }

    err = heif_context_set_primary_image(ctx, grid_handle);
    if (err.code != heif_error_Ok) {
        fprintf(stderr, "set_primary_image failed: %s\n", err.message);
        return 1;
    }

    for (int i = 0; i < 4; i++) {
        err = heif_context_add_image_tile(ctx, grid_handle, 0, i, tiles[i], encoder);
        if (err.code != heif_error_Ok) {
            fprintf(stderr, "add_image_tile failed at row %d: %s\n", i, err.message);
            return 1;
        }
    }

    err = heif_context_write_to_file(ctx, out_file);
    if (err.code != heif_error_Ok) {
        fprintf(stderr, "write failed: %s\n", err.message);
        return 1;
    }

    printf("Generated %s\n", out_file);

    if (decode_mode) {
        heif_context* dec_ctx = heif_context_alloc();
        err = heif_context_read_from_file(dec_ctx, out_file, NULL);
        if (err.code != heif_error_Ok) {
            fprintf(stderr, "read failed: %s\n", err.message);
            return 1;
        }

        heif_image_handle* handle = NULL;
        err = heif_context_get_primary_image_handle(dec_ctx, &handle);
        if (err.code != heif_error_Ok) {
            fprintf(stderr, "get_primary failed: %s\n", err.message);
            return 1;
        }

        heif_image* img = NULL;
        err = heif_decode_image(handle, &img, heif_colorspace_YCbCr, heif_chroma_420, NULL);
        if (err.code != heif_error_Ok) {
            fprintf(stderr, "decode failed: %s\n", err.message);
            return 1;
        }
        printf("Decoded OK\n");
        heif_image_release(img);
        heif_image_handle_release(handle);
        heif_context_free(dec_ctx);
    }

    heif_image_handle_release(grid_handle);
    heif_encoder_release(encoder);
    for (int i = 0; i < 4; i++) heif_image_release(tiles[i]);
    heif_context_free(ctx);
    return 0;
}