/************************* 
 * Owen Sullivan
 * CPSC 2310-002
 * opsulli@clemson.edu
*************************/

#include "EncodeDecode.h"

int numDigits(int number) {
    
    // Zero can only ever have one digit
    if (number == 0) {

        return 1;

    }
    
    int numDigits = 0;

    // Loop and divide by ten each time to count the number of digits
    while (number != 0) {

        numDigits++;

        number = number / 10;

    }

    return numDigits;

}

void promptMenu(FILE *inFile, FILE *outFile) {

    // Present menu options and prompt for an input
    fprintf(stdout, "PPM Steganography Program:\n\n");
    fprintf(stdout, "What would you like to do?\n");
    fprintf(stdout, "\t1. Encode default (program-provided) message.\n");
    fprintf(stdout, "\t2. Provide your own message to encode.\n");
    fprintf(stdout, "\nYour choice (1 or 2): ");

    // Scan in user input
    int choice;
    fscanf(stdin, " %d", &choice);

    // Validate user input
    while (choice != 1 && choice != 2) {

        fprintf(stdout, "Invalid choice. Try again (1 or 2): ");
        fscanf(stdin, " %d", &choice);

    }

    // If user choice is option 1, encode preset message
    if (choice == 1) {

        char message[] = "The quick brown fox jumps over the lazy dog.";

        fprintf(stdout, "\n");

        encodeImage(inFile, message, strlen(message), outFile);

    // If user choice is option 2, prompt for and scan in the user's custom 
    // message
    } else {

        // Get the header from the input image
        header_t header = readHeader(inFile);
        // Reset the file pointer since encodeImage re-reads the header
        rewind(inFile);
        
        // Create message string with the maximum number of characters that the 
        // image could be encoded with
        char message[((header.width * header.height) / 3) + 1];

        fprintf(stdout, "\nEnter a message to encode (max %d characters): ", 
                ((header.width * header.height) / 3));
        
        // Create a dynamically-sized format specifier to truncate extra 
        // characters after maximum number of characters, with dynamic 
        // maximum number of characters for different image sizes
        char formatSpecifier[(8 + 
                             numDigits((header.width * header.height) / 3))];
        sprintf(formatSpecifier, " %%%d[^\n]", 
                ((header.width * header.height) / 3));

        fscanf(stdin, formatSpecifier, message);

        fprintf(stdout, "\nThank you!\n\n");

        // Encode custom message
        encodeImage(inFile, message, strlen(message), outFile);

    }

}

void encodeImage(FILE *inFile, char message[], int size, FILE *outFile) {

    // Get the header from the input image
    header_t header = readHeader(inFile);

    // Make sure the input image is of PPM type P6
    assert(strcmp(header.type, "P6") == 0);

    // Write the header, since there are no changes being made to the header
    writeHeader(outFile, header);

    // Create and allocate a pointer for the image's pixels
    pixel_t **image = allocateMemory(header.width, header.height);

    readImage(inFile, image, header.width, header.height);

    int j = 0;
    int k = 0;
    // Loop through the characters in the message array
    for (int i = 0; i < size; i++) {

        // Get the decimal of the next character to be encoded
        unsigned int decimal = message[i];

        // Create a pointer to act as an array for a 9-bit binary 
        // representation
        unsigned int *binary = decimalToBinary(decimal);

        // Loop through each bit in the 9-bit binary representation
        for (int l = 0; l < 9; l += 3) {

            unsigned int lastDigit;
            
            // Strip the last digit of the pixel value, then assign a new 
            // last digit from the binary array
            lastDigit = image[j][k].r % 10;
            image[j][k].r -= lastDigit;
            image[j][k].r += binary[l];

            // Strip the last digit of the pixel value, then assign a new 
            // last digit from the binary array
            lastDigit = image[j][k].g % 10;
            image[j][k].g -= lastDigit;
            image[j][k].g += binary[l + 1];

            // Strip the last digit of the pixel value, then assign a new 
            // last digit from the binary array
            lastDigit = image[j][k].b % 10;
            image[j][k].b -= lastDigit;
            image[j][k].b += binary[l + 2];

            // Move on to the next pixel in the row
            k++;

            // If the end of a row has been reached, move back to the 
            // beginning of a row and move to the next row
            if (k >= header.width) {

                k = 0;
                j++;

            }

        }

        // Free pointer
        free(binary);

    }

    // Write the encoded image to the output file
    writeImage(outFile, image, header.width, header.height);

    // Free image pointer memory
    freeMemory(image, header.height);

}

unsigned int *decimalToBinary(unsigned int decimal) {
    
    // Create a pointer to act as the binary array
    unsigned int *binary = malloc(9 * sizeof(unsigned int));

    // If the input decimal is 0, call helper function to set the contents 
    // of the binary array to all 0s
    if (decimal == 0) {

        clearBinaryArray(binary);

      // If the input decimal is above 0
    } else {

        // Call helper function to set the contents of the binary array 
        // to all 0s, just in case the value being converted is not 9 bits
        clearBinaryArray(binary);
        
        int i = 0;
        // Loop until decimal can no longer be divided by 2
        while (decimal > 0) {

            // Store the remainder of the decimal divided by 2 as the next bit
            binary[i++] = decimal % 2;
            decimal = decimal / 2;

        }

        // Reverse the order of the array so that the 9-bit binary is readable 
        // left-to-right
        for (i = 0; i < (9 / 2); i++) {

            unsigned int temp = binary[i];
            binary[i] = binary[9 - 1 - i];
            binary[9 - 1 - i] = temp;

        }

    }

    return binary;

}

void clearBinaryArray(unsigned int binary[]) {

    // Loop through 9 bits in binary array
    for (int i = 0; i < 9; i++) {

        // Use bitwise AND operator to set the current value to 0
        binary[i] = binary[i] & 0;

    }

}

unsigned int binaryToDecimal(unsigned int binary[]) {

    unsigned int decimal = 0;

    // Loop through 9 bits in binary array
    for (int i = 0; i < 9; i++) {

        // Formula for binary conversion:
        // Multiply value by 2, then add next bit
        decimal = (decimal * 2) + binary[i];

    }

    return decimal;

}

void decodePixels(header_t header, pixel_t **image, unsigned int binary[]) {

    int counter = 0;

    // Loop through rows of pixels
    for (int i = 0; i < header.height; i++) {
        
        // Loop through pixels in a given row
        for (int j = 0; j < header.width; j++) {

            // If the pixel value is valid (255 or less) and the last digit
            // of the pixel value is valid (1 or 0), add the last digit to 
            // binary array
            if (image[i][j].r <= 255 && image[i][j].r % 10 <= 1) {

                binary[counter++] = image[i][j].r % 10;

              // If the pixel valid is invalid (256 or more) or the last digit
              // of the pixel value is invalid (2 or more), return early
            } else {

                return;

            }
            // If the pixel value is valid (255 or less) and the last digit
            // of the pixel value is valid (1 or 0), add the last digit to 
            // binary array
            if (image[i][j].g <= 255 && image[i][j].g % 10 <= 1) {

                binary[counter++] = image[i][j].g % 10;

              // If the pixel valid is invalid (256 or more) or the last digit
              // of the pixel value is invalid (2 or more), return early
            } else {

                return;

            }
            // If the pixel value is valid (255 or less) and the last digit
            // of the pixel value is valid (1 or 0), add the last digit to 
            // binary array
            if (image[i][j].b <= 255 && image[i][j].b % 10 <= 1) {

                binary[counter++] = image[i][j].b % 10;

              // If the pixel valid is invalid (256 or more) or the last digit
              // of the pixel value is invalid (2 or more), return early
            } else {

                return;

            }

            // If 9 bits have been filled in binary array
            if (counter == 9) {

                // If the decimal equivalent of the 9-bit binary 
                // representation is valid ASCII (divisible by 128 with a 
                // remainder)
                if (binaryToDecimal(binary) % 128 != 0) {

                    // Print the character
                    fprintf(stdout, "%c", binaryToDecimal(binary) % 128);

                }
                
                // Reset counter for binary array
                counter = 0;
                // Set all values to 0 with helper function
                clearBinaryArray(binary);

            }

        }

    }

}

void decodeImage(FILE *inFile) {
    
    // Get the header from the input image
    header_t header = readHeader(inFile);

    // Create and allocate a pointer for the image's pixels
    pixel_t **image = allocateMemory(header.width, header.height);

    readImage(inFile, image, header.width, header.height);

    // Create a 9-bit array signifying a binary sequence, setting all bits to 0
    // to start
    unsigned int binary[9] = {0, 0, 0, 0, 0, 0, 0, 0, 0};

    fprintf(stdout, "Decoded message:\n");
    // Call function which finds encoded characters and prints them
    // (needed to be able to return early and break nested loops)
    decodePixels(header, image, binary);
    // Print a new line after the decoded characters have been printed
    fprintf(stdout, "\n");

    // Free image pointer memory
    freeMemory(image, header.height);

}