#include <stdatomic.h>
#include <stdint.h>
#include <stddef.h>
#include <stdio.h>
#include <stdlib.h>
#include <endian.h>
#include <string.h>

// Tool to combine multiple JPEG files into a "regressive" image that changes to something
// else as it loads. The starting images must be progressive JPEGs
// 
// IMPORTANT: All images must have the same frame header, since only one is allowed
// per file. This means: same resolution, same bit depth, same channel count, same
// subsampling and quantization table identifiers.
//
// Since quantization table layout is usualy not user visible, I recomend encoding all
// frames with the same software.
//
// The actual quantization tables and huffman codes do not need to be constant. 

// Uncomment to trim out all but the first scan in an image.
//#define FORCE_DC 1

// Size of filler data between frames, in 10s of kb
// Set to zero to disable
#define N_FILLER 0

// Video frames. Output is written to "out.jpg"
#define N_FRAMES (2)
char* frames[N_FRAMES] = {
//    "src1.jpg",
//    "src2.jpg"
    #include "frames.h"
};

FILE* src = NULL;
FILE* output = NULL;

// JPEG datastream primitives
void write_null_marker(uint8_t magic) {
    uint8_t b = 0xff;
    fwrite(&b, sizeof(uint8_t), 1, output);
    fwrite(&magic, sizeof(uint8_t), 1, output);
}
void write_marker(uint8_t magic, uint16_t length, void* data) {
    write_null_marker(magic);
    uint16_t length_field = htobe16(length + 2);
    fwrite(&length_field, sizeof(uint16_t), 1, output);
    fwrite(data, sizeof(uint8_t), length, output);
}

// Sequencing
int frame = 0;
void next_frame() {
    printf("SWITCHING!\n");
    if (frame >= N_FRAMES) {
        src = NULL;
    }

    // Write dummy data
    char* filler = malloc(10000);
    for (int i = 0; i < 10000; i++) {
        filler[i] = i*i + i;
    }
    for (int i = 0; i < N_FILLER; i++) write_marker(0xFE, 10000, filler);

    src = fopen(frames[frame], "r");
    frame++;
}

int main() {
    output = fopen("out.jpg", "w");

    // Write start of image
    write_null_marker(0xD8);

    next_frame();

    // Witty comment
    char* msg = "Handcrafted from the finest bits";
    write_marker(0xFE, strlen(msg), msg);

    uint8_t has_prefix = 0;

    while (src) {
        // For each marker...
        read_marker:
        uint8_t magic;
        int count = fread(&magic, sizeof(uint8_t), 1, src);
        if (count == 0) {
            if (frame >= N_FRAMES) {
                break;
            } else {
                next_frame();
            }
            printf("Warning: got EOF instead of marker!\n");
        }

        if (magic == 0xFF) {
            has_prefix = 1;
            continue;
        } else {
            if (!has_prefix) {
                printf("Missing 0xFF marker prefix.\n");
                printf("This does not apear to be a valid jpeg.\n");
            }
            has_prefix = 0;
        }
        
        handle_marker:
        printf("Marker %x\n", magic);
        if (magic == 0xD8) {
            // Start of image, ignore
        } else if (magic == 0xD9) {
            // End of image, switch to the other file
            has_prefix = 0;
            next_frame();
            if (src == NULL) {
                break;
            }
        } else if (magic == 0xC0) {
            // Start of frame (baseline): abort
            printf("This does not look like a progressive JPEG. Aborting.\n");
            return 1;
        } else {
            
            // Some varable size marker...
            uint16_t size;
            fread(&size, sizeof(uint16_t), 1, src);
            size = be16toh(size) - 2;

            uint8_t data[size];
            fread(data, sizeof(uint8_t), size, src);

            // Only copy the first image's start of frame marker
            if (magic == 0xC2 && frame != 1) continue;
            // Remove comments
            if (magic == 0xFE) continue;

            // Pass through...
            write_marker(magic, size, data);
            
            if (magic == 0xDA) {
                printf("[Image data]\n");
                printf("\tChannels: %d\n", data[0]);
                for (int i = 0; i < data[0]; i++) {
                    printf("\t... #%d is %d\n", i, data[1 + 2*i]);
                }
                printf("\tSpectral start: %d\n", data[size-3]);
                printf("\tSpectral end: %d\n", data[size-2]);
                printf("\tPOINT TRANSFORM: %x\n", data[size-1]);
                // Start of scan markers have trailing image data before the next marker
                uint8_t buffer = 0;
                int count = 1;
                while (count) {
                    count = fread(&buffer, sizeof(uint8_t), 1, src);
                    // Detect markers
                    if (buffer == 0xFF) {
                        fread(&magic, sizeof(uint8_t), 1, src);
                        if (magic == 0x00 || (magic >= 0xD0 && magic <= 0xD7)) {
                            // 0xFF 0x00 is an excaped 0xff byte.
                            // 0xFF 0xD0-0xD7 are restart markers which don't end the image data.
                            fwrite(&buffer, sizeof(uint8_t), 1, output);
                            fwrite(&magic, sizeof(uint8_t), 1, output);
                        } else {
                            // otherwise, it's a real marker.
                            has_prefix = 1;
                            #ifdef FORCE_DC
                            next_frame();
                            if (!src) break;
                            goto read_marker;
                            #endif
                            goto handle_marker;
                        }
                    } else {
                        fwrite(&buffer, sizeof(uint8_t), 1, output);                    
                    }
                }
            }
        }
    }

    // End of image
    write_null_marker(0xD9);
    
    return 0;
}