#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdbool.h>

// Globally-defined constants as relevant to FAT12
#define SECTOR_SIZE 512
#define DIRECTORY_ENTRY_SIZE 32
#define CLUSTER_OFFSET 2
#define BOOT_SECTOR 0
#define FAT_START 1
#define FAT_END 9
#define ROOT_DIRECTORY_START 19
#define ROOT_DIRECTORY_END 32
#define DATA_START 33
#define CLUSTER_INDEX_SIZE 3
#define NUM_FAT_CLUSTERS ((FAT_END * SECTOR_SIZE) / CLUSTER_INDEX_SIZE)
#define FIRST_MASK 0xF0
#define SECOND_MASK 0x0F
#define DIRECTORY_MASK 0x10
#define DELETED_FILE_PREFIX 0xE5
#define ZERO_FAT_ENTRY 0x000
#define FAT_EOF 0xFFF

// Variable to keep track of how many files have been recovered
int numFilesRecovered = BOOT_SECTOR;

// Function that returns the size of the entire image being used 
// to recover files
int getFileSize(FILE *image) {

    // Seek to the end of the file
    fseek(image, BOOT_SECTOR, SEEK_END);

    // Store the current state of the file pointer
    int size = ftell(image);

    // Seek back to the front of the file
    fseek(image, BOOT_SECTOR, SEEK_SET);

    return size;

}

// Function to store each sector of the image being used to recover files 
// as a separate array entry, so that it is more easily traversable
void splitImageSectors(FILE *image, int size, void *imageSectors[]) {

    // Loop through each sector by calculating the number of sectors
    for (int i = BOOT_SECTOR; i < (size / SECTOR_SIZE); i++) {

        // Allocate space for the pointer to the current sector
        imageSectors[i] = malloc(SECTOR_SIZE);

        // Read the raw data from the image
        fread(imageSectors[i], SECTOR_SIZE, FAT_START, image);

    }

}

// Function that stores the cluster indexes pointed to by each entry in 
// the FAT
void getClusterIndexes(void *imageSectors[], short clusters[]) {

    // Create a pointer (of type unsigned char because it's 1 byte) to store 
    // the entire FAT, allocated with (9 * 512)
    unsigned char *FAT = malloc(FAT_END * SECTOR_SIZE);
    // Loop from the start of the FAT (sector 1) to the end of the FAT 
    // (sector 9)
    for (int i = FAT_START; i <= FAT_END; i++) {

        // Loop through each byte in the current sector
        for (int j = BOOT_SECTOR; j < SECTOR_SIZE; j++) {

            // Copy the current byte to the FAT, incrementing the FAT pointer 
            // in the process so that the next byte is stored correctly
            memcpy(FAT++, imageSectors[i] + j, FAT_START);

        }

    }

    // Move the FAT pointer back to the beginning of the first sector
    FAT -= (FAT_END * SECTOR_SIZE);

    // Loop from the start of the FAT's first sector to the end of the FAT's 
    // last sector
    for (int i = BOOT_SECTOR; i < NUM_FAT_CLUSTERS; i++) {

        // Store the next three bytes in the FAT pointer, incrementing the 
        // pointer each time
        unsigned char byte1 = *FAT++;
        unsigned char byte2 = *FAT++;
        unsigned char byte3 = *FAT++;

        // Store the first cluster as the second half of the second byte 
        // added to the beginning of the first byte, incrementing i in the 
        // process
        clusters[i++] = (byte2 & SECOND_MASK) + byte1;
        // Reverse the third byte
        byte3 = ((byte3 & SECOND_MASK) << (CLUSTER_OFFSET * CLUSTER_OFFSET) | 
                 (byte3 & SECOND_MASK) >> (CLUSTER_OFFSET * CLUSTER_OFFSET));
        // Store the second cluster as the first half of the second byte 
        // added to the end of the third byte
        clusters[i] = (byte3 & SECOND_MASK) + ((byte3 & FIRST_MASK) | 
                       (byte2 >> (CLUSTER_OFFSET * CLUSTER_OFFSET)));

    }

}

// Function that removes characters in a string beyond a maximum size
char *stripExtraCharacters(char *string, int maxSize) {

    // If the string passed as a parameter is NULL, return NULL
    if (string == NULL) {

        return NULL;

    }
    
    // Create a pointer for the new string with the maximum size plus 1 (for 
    // null termination)
    char *newString = malloc(maxSize + FAT_START);

    // Copy the contents of the old string into the new string, up to the 
    // maximum size
    strncpy(newString, string, maxSize);
    // Null-terminate new string
    newString[maxSize] = '\0';

    return newString;

}

// Function to calculate the physical sector for a given logical sector
int getPhysicalSectorNumber(int logicalSectorNumber) {

    // Evaluate (33 + logicalSectorNumber - 2)
    return (DATA_START + logicalSectorNumber - CLUSTER_OFFSET);

}

// Function that prints information about each file found in the FAT image
void printFileEntry(char *file, bool isDeletedFile, int size) {

    // If isDeletedFile is true, indicating that the file is marked as 
    // deleted, print with "DELETED"
    if (isDeletedFile) {

        fprintf(stdout, "FILE\tDELETED\t%s\t%d\n", file, size);

        // Otherwise, print with "NORMAL"
    } else {

        fprintf(stdout, "FILE\tNORMAL\t%s\t%d\n", file, size);

    }

}

// Function that returns the number of digits in an integer
int getNumDigits(int value) {

    // Create a variable to keep track of the number of digits
    int numDigits = BOOT_SECTOR;

    // Loop until the value equals 0
    while (value != BOOT_SECTOR) {

        // Divide the value by 10
        value /= (FAT_END + FAT_START);

        // Increment numDigits
        numDigits++;

    }

    return numDigits;

}

// Function that recovers a file from the FAT image by looking through the 
// FAT for the next cluster and recovering data by sectors as appropriate
void recoverFile(void *imageSectors[], 
                 short clusters[], 
                 char *outputDirectory, 
                 bool isDeletedFile, 
                 char *fileExtension, 
                 short *firstLogicalSector, 
                 int fileSize) {
    
    // Create a pointer to store the full output file path
    char *outputFilePath;

    // If the file extension exists (is not empty)
    if (strcmp(fileExtension, "")) {

        // Allocate the outputFilePath pointer with size equal to the length 
        // of the output directory (with a trailing slash), the length of the 
        // word "file", the number of digits in the counter for number of files 
        // recovered, a "." before the extension, and the length of the file 
        // extension
        outputFilePath = malloc(strlen(outputDirectory) + 
                                strlen("/") + 
                                strlen("file") + 
                                getNumDigits(numFilesRecovered) + 
                                strlen(".") + 
                                strlen(fileExtension));

        // Create a formatted string with the output directory (with a trailing 
        // slash), the word "file" (as instructed for this program), the running 
        // file number (incremented in the process), a "." before the 
        // extension, and the file extension
        sprintf(outputFilePath, 
                "%s%s%s%d%s%s", 
                outputDirectory, 
                "/", 
                "file", 
                numFilesRecovered++, 
                ".", 
                fileExtension);

        // If the file extension does not exist (is empty)
    } else {

        // Allocate the outputFilePath pointer with size equal to the length 
        // of the output directory (with a trailing slash), the length of the 
        // word "file", and the number of digits in the counter for number of 
        // files recovered
        outputFilePath = malloc(strlen(outputDirectory) + 
                                strlen("/") + 
                                strlen("file") + 
                                getNumDigits(numFilesRecovered));

        // Create a formatted string with the output directory (with a trailing 
        // slash), the word "file" (as instructed for this program), and the 
        // running file number (incrementing in the process)
        sprintf(outputFilePath, 
                "%s%s%s%d", 
                outputDirectory, 
                "/", 
                "file", 
                numFilesRecovered++);

    }
    
    // Open a FILE pointer with the full output file path (in write mode)
    FILE *outputFile = fopen(outputFilePath, "wb");

    // Store the first cluster number
    int nextCluster = *firstLogicalSector;

    // Loop while the remaining file size is greater than 0
    while (fileSize > BOOT_SECTOR) {

        // If the file being recovered is marked as a deleted file and the 
        // FAT entry is not 0x000, break the loop and truncate the file early
        if (isDeletedFile && clusters[nextCluster] != ZERO_FAT_ENTRY) {

            break;

        }
        
        // Initialize the size of the data to write as 512
        int sizeToWrite = SECTOR_SIZE;
        // If the remaining file size is less than 512, use the remaining 
        // file size as the size to write
        if (fileSize / SECTOR_SIZE == BOOT_SECTOR) {

            sizeToWrite = fileSize;

        }

        // Write the bytes in the physical sector corresponding to the next 
        // cluster the file occupies (using helper function 
        // getPhysicalSectorNumber()) to the output file
        fwrite(imageSectors[getPhysicalSectorNumber(nextCluster)], 
               sizeToWrite, 
               FAT_START, 
               outputFile);

        // Decrement the remaining file size by 512
        fileSize -= SECTOR_SIZE;

        // If the file is marked as deleted, simply move to the next cluster 
        // for recovery
        if (isDeletedFile) {

            nextCluster++;

            // If the file is not marked as deleted, store the next cluster 
            // that the file occupies according to the FAT
        } else {

            nextCluster = clusters[nextCluster];

        }

    }

    // Close the output file
    fclose(outputFile);

}

// Function that recursively finds files in directories, determines their 
// names, attributes, and sizes, and sends that information to printFileEntry()
// This function is intended to be used with the root directory as a starting 
// point
void findFilesInDirectory(void *imageSectors[], 
                          short clusters[], 
                          unsigned char *directory, 
                          char *lastDirectoryPath, 
                          char *outputDirectory) {
    
    // Create a variable to store a count of the number of directory entries 
    // processed, used in case directories overflow across multiple sectors
    int numDirectoryEntries = BOOT_SECTOR;

    // Create a pointer to store the first logical sector of a directory pointer
    short *currentDirectoryFLS = malloc(CLUSTER_OFFSET);
    
    // Loop through each directory entry in the current directory
    for (int i = BOOT_SECTOR; i < (SECTOR_SIZE / DIRECTORY_ENTRY_SIZE); i++) {

        // Create a pointer for the file name
        // Calls helper function stripExtraCharacters(), passing the 
        // return value of strtok() from the start of the directory entry 
        // and the maximum size of a file name in FAT12
        char *fileName = stripExtraCharacters(strtok((char *)directory, " "), 
                                              (FAT_END - FAT_START));

        // If the file name is NULL, we know we're done looking through 
        // directory entries
        if (fileName == NULL) {

            break;

        }

        // Increment numDirectoryEntries
        numDirectoryEntries++;

        // Create a pointer for the file extension
        // Calls helper function stripExtraCharacters(), passing the 
        // return value of strtok() from the start of the file extension bytes 
        // in the directory entry and the maximum size of a file extension in 
        // FAT12
        char *fileExtension = stripExtraCharacters(strtok((char *)directory + 
                                                          (FAT_END - FAT_START), 
                                                          " "), 
                                                   CLUSTER_INDEX_SIZE);

        // Create a pointer to hold first logical sector data
        short *firstLogicalSector = malloc(CLUSTER_OFFSET);
        // Copy the first logical sector bytes in the directory entry 
        // to the firstLogicalSector pointer
        memcpy(firstLogicalSector, 
               (directory + 
                (DIRECTORY_ENTRY_SIZE - 
                 (CLUSTER_OFFSET * CLUSTER_INDEX_SIZE))), 
               CLUSTER_OFFSET);
        
        // If the attributes for the current directory entry include a 
        // set directory bit
        if (*(directory + (FAT_END - FAT_START) + CLUSTER_INDEX_SIZE) & 
            DIRECTORY_MASK) {

            // Create a pointer for the full file path, allocated with 
            // size equal to the length of the last directory path, the length 
            // of the directory name, plus 1 for a trailing slash
            char *filePath = malloc(strlen(lastDirectoryPath) + 
                                    strlen(fileName) + 
                                    FAT_START);
            // Concatenate the last directory path, directory name, and a 
            // trailing slash
            sprintf(filePath, "%s%s%s", lastDirectoryPath, fileName, "/");

            // If the directory name is not "." or ".." (pointers to current 
            // and previous directories)
            if (strcmp(fileName, ".") && strcmp(fileName, "..")) {
                
                // Recursively call findFilesInDirectory with the calculated 
                // physical sector (helper function getPhysicalSectorNumber()) 
                // and the sub-directory's path
                findFilesInDirectory(imageSectors, 
                                     clusters, 
                    imageSectors[getPhysicalSectorNumber(*firstLogicalSector)], 
                                     filePath, 
                                     outputDirectory);
            
                // If the directory name is ".", store the first logical sector 
                // so that sector overflow can be handled properly
            } else if (!(strcmp(fileName, "."))) {

                // Copy the first logical sector bytes in the directory entry 
                // to the currentDirectoryFLS pointer
                memcpy(currentDirectoryFLS, 
                       (directory + 
                        (DIRECTORY_ENTRY_SIZE - 
                         (CLUSTER_OFFSET * CLUSTER_INDEX_SIZE))), 
                       CLUSTER_OFFSET);

            }

            // If the attributes for the current directory entry do not have a 
            // set directory bit
        } else {

            // Create a pointer for the full file path
            char *filePath;

            // If the file extension exists (is not empty)
            if (strcmp(fileExtension, "")) {

                // Allocate the filePath pointer with size equal to the length 
                // of the last directory path, the length of the file 
                // name, plus 1 for a '.', and the length of the file extension
                filePath = malloc(strlen(lastDirectoryPath) + 
                                  strlen(fileName) + 
                                  FAT_START + 
                                  strlen(fileExtension));

                // Concatenate the last directory path, file 
                // name, a '.', and the file extension
                sprintf(filePath, 
                        "%s%s%s%s", 
                        lastDirectoryPath, 
                        fileName, 
                        ".", 
                        fileExtension);

                // If the file extension does not exist (is empty)
            } else {

                // Allocate the filePath pointer with size equal to the length 
                // of the last directory path and the length of the file name
                filePath = malloc(strlen(lastDirectoryPath) + strlen(fileName));

                // Concatenate the last directory path and the file name
                sprintf(filePath, "%s%s", lastDirectoryPath, fileName);

            }
            
            // Create a pointer to store the file size data
            int *fileSize = malloc(CLUSTER_OFFSET * CLUSTER_OFFSET);
            // Copy the file size data to the fileSize pointer
            memcpy(fileSize, 
                   (directory + 
                    (DIRECTORY_ENTRY_SIZE - (CLUSTER_OFFSET * CLUSTER_OFFSET))), 
                   (CLUSTER_OFFSET * CLUSTER_OFFSET));

            // Create a variable to store the status of whether the file being 
            // recovered is marked as deleted
            bool isDeletedFile = false;
            
            // Store the first bit in the current directory entry
            unsigned char firstNameBit = *directory;
            // If the first bit in the file name matches 0xE5
            if (firstNameBit == DELETED_FILE_PREFIX) {

                // Change the character represented by 0xE5 to a printable 
                // character
                filePath[strlen(lastDirectoryPath)] = '_';

                // Update isDeletedFile status
                isDeletedFile = true;

            }

            // Call helper function printFileEntry() with the full file 
            // path, isDeletedFile for the second argument to indicate the 
            // file status, and the file size
            printFileEntry(filePath, isDeletedFile, *fileSize);

            // Call helper function recoverFile(), passing every piece of 
            // relevant information (there's a lot) to create a FILE output 
            // and navigate through the FAT image
            recoverFile(imageSectors, 
                        clusters, 
                        outputDirectory, 
                        isDeletedFile, 
                        fileExtension, 
                        firstLogicalSector, 
                        *fileSize);

        }

        // Advance the directory entry pointer
        directory += DIRECTORY_ENTRY_SIZE;

    }

    // If 16 directory entries have been read and the FAT entry for the current 
    // cluster does not point to EOF, recursively call findFilesInDirectory() 
    // to look for files in sub-directories that span across multiple sectors
    if (numDirectoryEntries == (DIRECTORY_ENTRY_SIZE / CLUSTER_OFFSET) && 
        clusters[*currentDirectoryFLS] != FAT_EOF) {

        findFilesInDirectory(imageSectors, 
                             clusters, 
        imageSectors[getPhysicalSectorNumber(clusters[*currentDirectoryFLS])], 
                             lastDirectoryPath, 
                             outputDirectory);

    }

}

// Function to find all files in the FAT image, starting from the root directory
void getFiles(void *imageSectors[], short clusters[], char *outputDirectory) {

    // Create a pointer (of type unsigned char because it's 1 byte) to store 
    // the entire root directory, allocated with (14 * 512)
    unsigned char *rootDirectory = malloc((DATA_START - ROOT_DIRECTORY_START) * 
                                           SECTOR_SIZE);
    // Loop from the start of the root directory (sector 19) to the end of 
    // the root directory (sector 32)
    for (int i = ROOT_DIRECTORY_START; i <= ROOT_DIRECTORY_END; i++) {

        // Loop through each byte in the current sector
        for (int j = BOOT_SECTOR; j < SECTOR_SIZE; j++) {

            // Copy the current byte to the root directory, incrementing the 
            // root directory pointer in the process so that the next byte is 
            // stored correctly
            memcpy(rootDirectory++, imageSectors[i] + j, FAT_START);

        }

    }

    // Move the root directory pointer back to the beginning of the first sector
    rootDirectory -= ((DATA_START - ROOT_DIRECTORY_START) * SECTOR_SIZE);

    // Call helper function findFilesInDirectory to start finding files 
    // recursively (handles printing indirectly)
    findFilesInDirectory(imageSectors, 
                         clusters, 
                         rootDirectory, 
                         "/", 
                         outputDirectory);

}

// Main function
int main(int argc, char **argv) {

    // Open the image provided as a command-line argument
    FILE *image = fopen(argv[FAT_START], "rb");

    // Get the size of the image file by calling helper function getFileSize()
    int size = getFileSize(image);

    // Create an array of pointers that will store raw data
    void *imageSectors[size / SECTOR_SIZE];
    // Call helper function splitImageSectors() to divide raw image data 
    // into sectors
    splitImageSectors(image, size, imageSectors);

    // Close the file pointer to the image, since sectors are now stored
    fclose(image);

    // Create an array to store the cluster indexes
    short clusters[NUM_FAT_CLUSTERS];
    // Call helper function getClusterIndexes() to divide up the FAT cluster 
    // indexes
    getClusterIndexes(imageSectors, clusters);

    // Call helper function getFiles() to start searching through the FAT image 
    // for files
    getFiles(imageSectors, clusters, argv[CLUSTER_OFFSET]);

}