#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <MagickWand/MagickWand.h>

int isPowerOfTwo(int x) {return x && (!(x & (x - 1)));}

unsigned int next_power_of_2(unsigned int n) {
    if (n == 0) {
        return 1;
    }
    unsigned int p = 1;
    while (p < n) {
        p <<= 1;
    }
    return p;
}

unsigned int nearest_power_of_2(unsigned int n) {
    if (n == 0) {
        return 1;
    }
    unsigned int next_power = next_power_of_2(n);
    unsigned int prev_power = next_power >> 1;

    if (n - prev_power <= next_power - n) {
        return prev_power;
    } else {
        return next_power;
    }
}

struct options {
    uint32_t r;
    int s;
    uint8_t c;
    int o;
    int l;
    int w;
    int h;
};

void printHelp() {
    printf(
        "usage: kaveat <target MOV> <source images> [options]\n"
        "Where each source begins with +, and is added as an animation frame.\n"
        "A maximum of 360 different frames can be used for a single avatar.\n"
        "If the size of any frame does not match, it will be resized to match the first frame.\n"
        "\n"
        "Available options:\n"
        "   -r #    Maximum resolution, where # is a power of 2. For example, 256 pixels would be \"-r 8\"\n"
        "   -s      Stretch the image when non-pot or mismatched size, as opposed to fitting it within.\n"
        "   -c #    Defines the number of colors in the palette, counting from 1. Maximum 256.\n"
        "   -o      Opaque mode. Do not make the first color in the image palette transparent.\n"
        "               The background of letterboxed images is internally set to the first palette color,\n"
        "               so using this setting with different image sizes will cause visible bars to appear\n"
        "               unless you combine it with the -s option, which will always fill the canvas for each\n"
        "               frame. The same is true for all imported frames which contain an alpha channel.\n"
//        "   -l          Enable Lanczos filtering when scaling the image.\n"  //point filter unimplemented
        //"   -w      Specify the number of subdivisions to make horizontally per frame.\n"
        //"               Use this if you are importing a spritesheet. Each cut is perfectly even.\n"
        //"               Note that this process happens AFTER importing each image, and will be\n"
        //"               applied to each one separately. DO NOT use this to combine a spritesheet"
        //"               with individual frames or you will have a very, very bad time."
        //"   -h      Same as -w, but vertical. If used with -w, KAvEAT will create each chunk\n"
        //"               from each row left-to-right first, not each column top-to-bottom.\n"
    );
}

void printNoFrames() { printf("You need to specify at least one frame to use, goober.\n"); }

int main(const int argc, char *argv[]) {
    printf(
        "Knowledge Adventure Worlds Easy Avatar Tool, version 1\n"
           "Copyright (c) 2025 Brett \"bonkmaykr\" Bergstrom\n"
           "This program is free software under the MIT license.\n"
           "\n"
    );
    if (argc < 2) { printHelp(); return 0; }
    // Strings allocated with variable size according to filename need to be done on the heap due to an MSVC bug.
    char* filename = (char*)malloc(strlen(argv[1])*sizeof(char)); // MSVC
    strcpy(filename, argv[1]);
    printf("Creating avatar %s.mov\n", filename);
    if (argc < 3) { printNoFrames(); return 0; }

    MagickWand* mw = NULL;
    MagickWandGenesis();
    mw = NewMagickWand();

    struct options opt;
    opt.r = 0;
    opt.s = 0;
    opt.c = 255;
    opt.o = 0;
    opt.l = 0;
    opt.w = 0;
    opt.h = 0;
    uint16_t numFrames = 0;
    char* frames[360];

    int i = 2;
    while (i < argc) {
        if (argv[i][0] == '+') {
            frames[numFrames] = (char*)malloc(strlen(argv[i]) + 1);
            memcpy(frames[numFrames], &argv[i][1], strlen(argv[i]) + 1);
            printf("Adding frame to search index: %s\n", frames[numFrames]);

            numFrames++;
        } else if (strcmp(argv[i], "-r") == 0) {
            if (i == argc-1) { printf("No parameter for option: -r\n"); return 0; }
            int r = pow(2, trunc(atof(argv[i+1])));
            if (r > 4096) { printf("I really, really hope your image is not %d pixels in size.\n", r); }
            else if (r < 2) {printf("I'm sorry, Dave, I'm afraid I can't do that. (%d pixel(s) is too small, clamping to 2.)\n", r); r = 2;}
            printf("Limiting resolution per frame to %dx%d pixels.\n", r, r);
            opt.r = r;
            i++;
        } else if (strcmp(argv[i], "-s") == 0) {
            opt.s = !opt.s;
        } else if (strcmp(argv[i], "-o") == 0) {
            opt.o = !opt.o;
        } else if (strcmp(argv[i], "-c") == 0) {
            if (i == argc-1) { printf("No parameter for option: -c\n"); return 0; }
            const int c = trunc(atof(argv[i+1]));
            if (c < 1 || c > 256) { printf("Invalid number of colors: %s\n", argv[i+1]); return 0; }
            opt.c = c - 1;
            i++;
        } else {
            printf("I don't know what \"%s\" means!\n", argv[i]);
            return 0;
        }
        i++;
    }
    if (numFrames == 0) { printNoFrames(); return 0; }

    if (MagickReadImage(mw, frames[0]) == MagickFalse) { printf("ERROR trying to read %s\n", frames[0]); return 0;}
    MagickSetImageFormat(mw, "BMP3");
    MagickSetImageDepth(mw, 8);
    MagickSetImageType(mw, PaletteType);
    MagickSetImageColorspace(mw, UndefinedColorspace);
    uint16_t width = MagickGetImageWidth(mw);
    uint16_t height = MagickGetImageHeight(mw);

    // -r option
    if (opt.r > 0) {
        if (opt.s > 0) {
            printf("Squishing to resolution limit...\n");
            if (width > opt.r) {
                MagickResizeImage(mw, nearest_power_of_2(width), nearest_power_of_2(height), LanczosFilter);
                width = opt.r;
            }
            if (height > opt.r) {
                MagickResizeImage(mw, nearest_power_of_2(width), nearest_power_of_2(height), LanczosFilter);
                height = opt.r;
            }
        }
        else {
            printf("Fitting to resolution limit...\n");
            if (width > opt.r || height > opt.r) {
                if (width > height) {
                    const double ratio = (double)height/(double)width;
                    printf("Letterboxing image to %ux%u.\n", opt.r, opt.r*ratio);
                    PixelWand* pw = NewPixelWand();
                    MagickGetImagePixelColor(mw, 0, 0, pw);
                    MagickSetImageBackgroundColor(mw, pw);

                    if (ratio < 1) MagickResizeImage(mw, opt.r, opt.r*ratio, LanczosFilter);
                    //MagickExtentImage(mw, opt.r, opt.r, 0, (opt.r-(opt.r*ratio))*-0.5);

                    width = opt.r;
                    height = opt.r;
                } else if (height > width) {
                    const double ratio = (double)width / (double)height;
                    printf("Pillarboxing image to %ux%u.\n", opt.r*ratio, opt.r);
                    PixelWand* pw = NewPixelWand();
                    MagickGetImagePixelColor(mw, 0, 0, pw);
                    MagickSetImageBackgroundColor(mw, pw);

                    if (ratio < 1) MagickResizeImage(mw, opt.r*ratio, opt.r, LanczosFilter);
                    //MagickExtentImage(mw, opt.r, opt.r, (opt.r-(opt.r*ratio))*-0.5, 0);

                    width = opt.r;
                    height = opt.r;
                } else /*if (height == width)*/ {
                    printf("Resizing image to %ux%u.\n", opt.r, opt.r);
                    MagickResizeImage(mw, opt.r, opt.r, LanczosFilter);

                    width = opt.r;
                    height = opt.r;
                }
            }
        }
    }

    if (!isPowerOfTwo(width) || !isPowerOfTwo(height)) printf("Uh-oh! Image size %hux%hu is not a power of two!\n", width, height);
    if (opt.s > 0 || (!isPowerOfTwo(width) && !isPowerOfTwo(height))) { // both are bad!
        printf("Resizing image to %ux%u.\n", nearest_power_of_2(width), nearest_power_of_2(height));
        MagickResizeImage(mw, nearest_power_of_2(width), nearest_power_of_2(height), LanczosFilter);
    }
    else if (!isPowerOfTwo(width) && isPowerOfTwo(height)) { // just the width is bad
        printf("Letterboxing image to %ux%u.\n", height, height);
        PixelWand* pw = NewPixelWand();
        MagickGetImagePixelColor(mw, 0, 0, pw);
        MagickSetImageBackgroundColor(mw, pw);

        const double ratio = (double)height/(double)width;
        if (ratio < 1) MagickResizeImage(mw, height, height*ratio, LanczosFilter);
        MagickExtentImage(mw, height, height, 0, (height-(height*ratio))*-0.5);
    } else if (!isPowerOfTwo(height) && isPowerOfTwo(width)) { // just the height is bad
        printf("Pillarboxing image to %ux%u.\n", width, width);
        PixelWand* pw = NewPixelWand();
        MagickGetImagePixelColor(mw, 0, 0, pw);
        MagickSetImageBackgroundColor(mw, pw);

        const double ratio = (double)width / (double)height;
        if (ratio < 1) MagickResizeImage(mw, width*ratio, width, LanczosFilter);
        MagickExtentImage(mw, width, width, (width-(width*ratio))*-0.5, 0);
    }

    width = MagickGetImageWidth(mw);
    height = MagickGetImageHeight(mw);
    double ratio1 = (double)width / (double)height;
    double ratio2;
    PixelWand* pw = NewPixelWand();
    MagickGetImagePixelColor(mw, 0, 0, pw);

    printf("Importing %d frames...\n", numFrames);
    i = 1;
    while (i < numFrames) {
        // Letterbox background default to transparent pixel
        MagickSetImageBackgroundColor(mw, pw);

        // Load the image
        if (MagickReadImage(mw, frames[i]) == MagickFalse) { printf("ERROR trying to read %s\n", frames[i]); return 0;}
        MagickSetImageFormat(mw, "BMP3");
        MagickSetImageDepth(mw, 8);
        MagickSetImageType(mw, PaletteType);
        MagickSetImageColorspace(mw, UndefinedColorspace);

        // Resize the image to fit the first frame
        ratio2 = (double)MagickGetImageWidth(mw) / (double)MagickGetImageHeight(mw);
        MagickSetImageGravity(mw, CenterGravity);
        if (MagickGetImageHeight(mw) != height || MagickGetImageWidth(mw) != width) {
            if (opt.s == 0) {
                if (ratio1 > ratio2) {
                    double scale = (double)height/(double)MagickGetImageHeight(mw);
                    MagickResizeImage(mw, MagickGetImageWidth(mw)*scale, height, LanczosFilter);
                } else {
                    double scale = (double)width/(double)MagickGetImageWidth(mw);
                    MagickResizeImage(mw, width, MagickGetImageHeight(mw)*scale, LanczosFilter);
                }
            } else {
                if (MagickGetImageWidth(mw) != width || MagickGetImageHeight(mw) != height) {
                    double scalew = (double)width/(double)MagickGetImageWidth(mw);
                    double scaleh = (double)height/(double)MagickGetImageHeight(mw);
                    MagickResizeImage(mw, MagickGetImageWidth(mw)*scalew, MagickGetImageHeight(mw)*scaleh, LanczosFilter);
                }
            }
        }

        i++;
    }

    // Create spritesheet from frames
    MagickResetIterator(mw);
    mw = MagickAppendImages(mw, MagickFalse);

    int piecesX = width/256;   if (width < 256) piecesX = 1;
    int piecesY = height/256;  if (height < 256) piecesY = 1;
    printf("Cutting up frames into %dx%d chunks...\n", width/piecesX, height/piecesY);

    i = 0;
    remove("segments.txt");
    FILE* partlist = fopen("segments.txt", "a");
    while (i < numFrames) {
        int ih = 0;
        const int chunkw = width/piecesX;
        const int chunkh = height/piecesY;
        while (ih < piecesY) {
            int iw = 0;
            while (iw < piecesX) {
                MagickWand* chunk = CloneMagickWand(mw);
                MagickSetImageFormat(chunk, "BMP3");
                MagickSetImageDepth(chunk, 8);
                MagickSetImageType(chunk, PaletteType);
                MagickSetImageColorspace(chunk, UndefinedColorspace);
                MagickCropImage(chunk, chunkw, chunkh, (chunkw*iw)+(width*i), chunkh*ih);

                printf("Saving frame #%d, chunk offset %d,%d\n", i+1, iw, ih);
                char* framename = (char*)malloc(12 * piecesX * piecesY * sizeof(char)); // MSVC
                sprintf(framename, "f%d%d%d.bmp", i, iw, ih);
                fprintf(partlist, "f%d%d%d.bmp\n", i, iw, ih);

                MagickWriteImage(chunk, framename);

                iw++;
            }
            ih++;
        }
        i++;
    }
    fclose(partlist);

    // Pass on images to COMPIMG
    printf("Encoding texture file...\n");
    char* command = (char*)malloc((54 + 3 * abs(1-opt.o) + strlen(filename))*sizeof(char)); // MSVC
    char argTransparent[3];
    if (opt.o == 0) sprintf(argTransparent, " -t");
#ifdef _MSC_VER
    sprintf(command, "compimg.exe -ace -c%d -r0 -pt%s -emov -M%s +segments.txt", opt.c + 1, argTransparent, filename);
#else
    sprintf(command, "./compimg.exe -ace -c%d -r0 -pt%s -emov -M%s +segments.txt", opt.c+1, argTransparent, filename);
#endif
    printf("%s\n", command);
    if (system(command)) { printf("ERROR trying to run COMPIMG!\n"); return 0; }

    printf("\n\n!!! DONE !!!\nYour avatar name is: %s%ds*%dh*%dv*.mov\n", filename, numFrames, piecesX, piecesY);

	if(mw) mw = DestroyMagickWand(mw);
	MagickWandTerminus();
    // don't clean up heap strings since the OS will reclaim the memory anyway
    return 0;
}