//     _    _                      _______ _                 _              _       _             
//    | |  | |                    |__   __| |               | |            (_)     (_)            
//    | |__| | __ _ _ __  _ __  _   _| |  | |__   __ _ _ __ | | _____  __ _ ___   ___ _ __   __ _ 
//    |  __  |/ _` | '_ \| '_ \| | | | |  | '_ \ / _` | '_ \| |/ / __|/ _` | \ \ / / | '_ \ / _` |
//    | |  | | (_| | |_) | |_) | |_| | |  | | | | (_| | | | |   <\__ \ (_| | |\ V /| | | | | (_| |
//    |_|  |_|\__,_| .__/| .__/ \__, |_|  |_| |_|\__,_|_| |_|_|\_\___/\__, |_| \_/ |_|_| |_|\__, |
//                 | |   | |     __/ |                                 __/ |                 __/ |
//                 |_|   |_|    |___/                                 |___/                 |___/ 


#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <getopt.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <string.h>

#include "utils.h"
#include "fat32.h"

int main(int argc, char **argv) {

    /* Deal with optional arguments */
    int ch;
    char *disk = NULL;
    int iFlag = 0, lFlag = 0, rMode = 0;
    char *fileName = NULL;
    char *sCode = NULL;
    while ((ch = getopt(argc, argv, "ilr:R:s:")) != -1) {
        switch (ch) {
            case 'i': iFlag = 1; break;
            case 'l': lFlag = 1; break;
            case 'r': rMode = 1; fileName = optarg; break;
            case 'R': rMode = 2; fileName = optarg; break;
            case 's': sCode = optarg; break;
            default: usage();
        }
    }

    /* Check validity for fixed composition of optional arguments */
    if (rMode == 2 && sCode == NULL) usage();
    if (rMode == 0 && sCode != NULL) usage();

    /* Create memory map for the FAT32 disk file */
    if (argc - optind != 1) usage();
    int fdisk = open(argv[optind], O_RDWR);
    struct stat fstats;
    fstat(fdisk, &fstats);
    disk = mmap(NULL, fstats.st_size, PROT_READ | PROT_WRITE, MAP_SHARED, fdisk, 0);

    /* Read the boot entry */
    BootEntry *bootEntry = malloc(sizeof(BootEntry));
    memcpy(bootEntry, disk, 90);

    /* Read the root directory, may spread across multiple clusters */
    DirEntry **rootDirEntries = malloc(sizeof(DirEntry *) * 512);
    int dirEntryCount = 0;
    unsigned int rootDirCluster = bootEntry -> BPB_RootClus;
    while (rootDirCluster < FAT32EOF) {
        unsigned long posoff = followFAT(bootEntry, &rootDirCluster, disk, 1);
        /* Traverse directory entries in the current cluster of the root directory */
        for (int i = 0; i * 32 < bootEntry -> BPB_BytsPerSec * bootEntry -> BPB_SecPerClus; i ++) {
            rootDirEntries[dirEntryCount] = malloc(sizeof(DirEntry));
            memcpy(rootDirEntries[dirEntryCount], disk + posoff + i * 32, 32);
            if (rootDirEntries[dirEntryCount] -> DIR_Name[0] == 0x00) break;
            dirEntryCount ++;
        }
    }

    /* Deal with disk information flag */
    if (iFlag) printFsInfo(bootEntry);

    /* Deal with root directory information flag */
    if (lFlag) {
        int entryCount = 0;
        for (int i = 0; i < dirEntryCount; i ++) {
            if (printDirInfo(rootDirEntries[i]) == 0) entryCount ++;
        }
        printf("Total number of entries = %d\n", entryCount);
    }

    /* Deal with recovering contiguously allocated files */
    if (rMode == 1) {
        int target = delContEntryIndex(fileName, sCode, disk, bootEntry, rootDirEntries, dirEntryCount);
        if (target == -1) printf("%s: file not found\n", fileName);
        else if (target == -2) printf("%s: multiple candidates found\n", fileName);
        else {
            unsigned int i = 0;
            DirEntry *targetEntry = rootDirEntries[target];
            unsigned int targetCluster = (targetEntry -> DIR_FstClusHI << 16) | targetEntry -> DIR_FstClusLO;
            /* Check if the target file is empty */
            if (targetCluster != 0) {
                /* Update FATs for the contiguous next cluster */
                while (targetEntry -> DIR_FileSize > bootEntry -> BPB_BytsPerSec * bootEntry -> BPB_SecPerClus * (i + 1)) {
                    updateFATs(bootEntry, targetCluster + i, disk, targetCluster + i + 1);
                    i ++;
                }
                /* Update FATs for the last cluster */
                updateFATs(bootEntry, targetCluster + i, disk, FAT32EOF);
            }
            /* Recover the name of the deleted file */
            int targetInClus = target % (bootEntry -> BPB_BytsPerSec * bootEntry -> BPB_SecPerClus / 32);
            unsigned int targetDirClus = bootEntry -> BPB_RootClus;
            for (int j = 0; (j + 1) * (bootEntry -> BPB_BytsPerSec * bootEntry -> BPB_SecPerClus / 32) <= target; j ++) {
                followFAT(bootEntry, &targetDirClus, disk, 1);
            }
            memcpy(disk + followFAT(bootEntry, &targetDirClus, disk, 0) + targetInClus * 32, fileName, 1);
            if (sCode == NULL) printf("%s: successfully recovered\n", fileName);
            else printf("%s: successfully recovered with SHA-1\n", fileName);
        }
    }

    /* Deal with recovering non-contiguously allocated files */
    else if (rMode == 2) {
        unsigned int *targets = malloc(sizeof(unsigned int) * 6);
        int target = delNonContClusters(fileName, sCode, disk, bootEntry, rootDirEntries, dirEntryCount, targets);
        if (target == -1) printf("%s: file not found\n", fileName);
        else {
            /* Check if the target file is empty */
            if (targets[0] != 0) {
                /* Update FATs for the possibly non-contiguous next cluster */
                for (unsigned int i = 1; i < targets[0]; i ++) {
                    updateFATs(bootEntry, targets[i], disk, targets[i + 1]);
                }
                /* Update FATs for the last cluster */
                updateFATs(bootEntry, targets[targets[0]], disk, FAT32EOF);
            }
            /* Recover the name of the deleted file */
            int targetInClus = target % (bootEntry -> BPB_BytsPerSec * bootEntry -> BPB_SecPerClus / 32);
            unsigned int targetDirClus = bootEntry -> BPB_RootClus;
            for (int j = 0; (j + 1) * (bootEntry -> BPB_BytsPerSec * bootEntry -> BPB_SecPerClus / 32) <= target; j ++) {
                followFAT(bootEntry, &targetDirClus, disk, 1);
            }
            memcpy(disk + followFAT(bootEntry, &targetDirClus, disk, 0) + targetInClus * 32, fileName, 1);
            printf("%s: successfully recovered with SHA-1\n", fileName);
        }
    }

    /* Free allocated heap memory */
    free(bootEntry);
    for (int i = 0; i < dirEntryCount; i ++) free(rootDirEntries[i]);
    free(rootDirEntries);

    /* Flush changes and unmap the FAT32 disk file */
    msync(disk, fstats.st_size, MS_SYNC);
    munmap(disk, fstats.st_size);
    close(fdisk);
    return 0;

}