#include <stdio.h>
#include <stdint.h>
#include <string.h>

void encrypt (uint32_t* v, uint32_t* k) 
{
    uint32_t v0=v[0], v1=v[1], sum=0, i;           /* set up */
    uint32_t delta=0x9e3779b9;                     /* a key schedule constant */
    uint32_t k0=k[0], k1=k[1], k2=k[2], k3=k[3];   /* cache key */
    for (i=0; i < 32; i++) {                       /* basic cycle start */
        sum += delta;
        v0 += ((v1<<4) + k0) ^ (v1 + sum) ^ ((v1>>5) + k1);
        v1 += ((v0<<4) + k2) ^ (v0 + sum) ^ ((v0>>5) + k3);
    }                                              /* end cycle */
    v[0]=v0; v[1]=v1;
}

void decrypt (uint32_t* v, uint32_t* k) 
{
    uint32_t delta=0x9e3779b9;
    uint32_t v0=v[0], v1=v[1], sum=(delta << 5), i; 
    uint32_t k0=k[0], k1=k[1], k2=k[2], k3=k[3];

    for (i=0; i < 32; i++) {
        v1 -= ((v0<<4) + k2) ^ (v0 + sum) ^ ((v0>>5) + k3);
        v0 -= ((v1<<4) + k0) ^ (v1 + sum) ^ ((v1>>5) + k1);
        sum -= delta;
    }
    v[0]=v0; v[1]=v1;
}



int treatFile(FILE* input, FILE* output, uint32_t key[4], char method) {

    int n;
    uint32_t buffer[2];
    int m;

    //TODO: Padding management

    switch (method)
    {
    case 'd':
        while ((n = fread(buffer, 4, 2, input)) > 0) {
            decrypt(buffer, key);

            if (( m = fwrite(buffer, 4, 2, output)) != 2) {
                return 1;
            }
        }
        break;

    case 'e':
        while ((n = fread(buffer, 4, 2, input)) > 0) {
            encrypt(buffer, key);

            if (( m = fwrite(buffer, 4, 2, output)) != 2) {
                return 1;
            }
        }
        break;
    }

    return 0;
}

int main(int argc, char const *argv[])
{
    if (argc < 5) {
        printf("Usage: %s e/d <key> <input> <output>\n", argv[0]);
        return 0;
    }
    
    if (!strcmp(argv[1], "-d") && !strcmp(argv[1], "-e")) {
        printf("Invalid operation, use either e(ncrypt) or d(ecrypt).");
        return 1;
    }

    uint32_t key[4];
    FILE* keyFile = fopen(argv[2], "r");
    if (!keyFile) {
        printf("Unable to open key file.\n");
        return 1;
    }

    if (fread(key, 4, 4, keyFile) != 4) {
        printf("Invalid key length.\n");
        return 1;
    }

    fclose(keyFile);

    FILE* input = fopen(argv[3], "r");
    if (!input) {
        printf("Unable to open input file.\n");
        return 1;
    }

    FILE* output = fopen(argv[4], "w");
    if (!output) {
        printf("Unable to open output file.\n");
        return 1;
    }

    int failure = treatFile(input, output, key, argv[1][1]);

    fclose(input);
    fclose(output);

    if (failure) {
        printf("Error during processing.\n");
        return 2;
    }

    return 0;
}