#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <fcntl.h>
#include <stdint.h>
#include <assert.h>
#include <string.h>


#define SECTOR_HEAD(buf, num) (buf + (num * 512))

#define FAT1_START 1
#define FAT2_START (bpb_ptr) (bpb_ptr->FATSz16 + FAT1_START)
#define ROOTDIR_START (bpb_ptr) (bpb_ptr->FATSz16 + FAT2_START)

#define BPB_BytsPerSec 512

#define CLUSTER_BAD     0xff7
#define CLUSTER_FREE    0x000
#define CLUSTER_FINAL   0xfff
#define CLUSTER_RSRVD   0xff8

#define ATTR_READ_ONLY  0x01
#define ATTR_HIDDEN     0x02
#define ATTR_SYSTEM     0x04
#define ATTR_VOLUME_ID  0x08
#define ATTR_DIRECTORY  0x10
#define ATTR_ARCHIVE    0x20
#define ATTR_LONG_NAME  ATTR_READ_ONLY | ATTR_HIDDEN | ATTR_SYSTEM | ATTR_VOLUME_ID

typedef struct BPB_type
{
    uint8_t jmpBoot[3];
    uint8_t OEMName[8];
    uint16_t    BytsPerSec;
    uint8_t     SecPerClus;
    uint16_t    RsvdSecCnt;
    uint8_t     NumFATs;
    uint16_t    RootEntCnt;
    uint16_t    TotSec16;
    uint8_t     Media;
    uint16_t    FATSz16;
    uint16_t    SecPerTrk;
    uint16_t    NumHeads;
    uint32_t    HiddSec;
    uint32_t    TotSec32;

    uint8_t     DrvNum;
    uint8_t     Reserved1;
    uint8_t     BootSig;
    uint32_t    VollD;
    uint8_t VolLab[11];
    uint8_t FilSysType[8];
    uint8_t BootCode[448];
    uint16_t    Signature_word;
}__attribute__((packed)) BPB_type;

typedef struct DIR_type {
    uint8_t Name [11];
    uint8_t     Attr;
    uint8_t     NTRes;
    uint8_t     CrtTimeTenth;
    uint16_t    CrtTime;
    uint16_t    CrtDate;
    uint16_t    LstAccDate;
    uint16_t    FstClusHI;
    uint16_t    WrtTime;
    uint16_t    WrtDate;
    uint16_t    FstClusLO;
    uint32_t    FileSize;
} __attribute__((packed)) DIR_type;

char* img_buffer;
size_t img_size;
static void read_img(const char* filename) {
    FILE* fp = fopen(filename, "rb");
    fseek(fp, 0, SEEK_END);
    img_size = ftell(fp);
    printf("Reading %ld bytes from %s\n", img_size, filename);
    fseek(fp, 0, SEEK_SET);
    img_buffer = (char*)malloc(img_size);
    fread(img_buffer, img_size, 1, fp);
    fclose(fp);
}

static void print_bpb(BPB_type* bpb_ptr) {
    printf("jmpBoot:    0x%2x 0x%2x 0x%2x\n", bpb_ptr->jmpBoot[0], bpb_ptr->jmpBoot[1], bpb_ptr->jmpBoot[2]);
    printf("OEMName:    "); for (int i = 0; i < 8; ++ i) putchar(bpb_ptr->OEMName[i]); putchar('\n');
    printf("BytsPerSec: %u\n", bpb_ptr->BytsPerSec);
    printf("SecPerClus: %u\n", bpb_ptr->SecPerClus);
    printf("RsvdSecCnt: %u\n", bpb_ptr->RsvdSecCnt);
    printf("NumFATs:    %u\n", bpb_ptr->NumFATs);
    printf("RootEntCnt: %u\n", bpb_ptr->RootEntCnt);
    printf("TotSec16:   %u\n", bpb_ptr->TotSec16);
    printf("Media:      %u\n", bpb_ptr->Media);
    printf("FATSz16:    %u\n", bpb_ptr->FATSz16);
    printf("SecPerTrk:  %u\n", bpb_ptr->SecPerTrk);
    printf("NumHeads:   %u\n", bpb_ptr->NumHeads);
    printf("HiddSec:    %u\n", bpb_ptr->HiddSec);
    printf("TotSec32:   %u\n", bpb_ptr->TotSec32);
    printf("DrvNum:     %u\n", bpb_ptr->DrvNum);
    printf("Reserved1:  %u\n", bpb_ptr->Reserved1);
    printf("BootSig:    %u\n", bpb_ptr->BootSig);
    printf("VollD:      %u\n", bpb_ptr->VollD);
    printf("VolLab:     "); for (int i = 0; i < 11; ++ i) putchar(bpb_ptr->VolLab[i]); putchar('\n');
    printf("FilSysType: "); for (int i = 0; i < 8; ++ i) putchar(bpb_ptr->FilSysType[i]); putchar('\n');
    printf("Sig:        0x%x\n", bpb_ptr->Signature_word);
}

static inline int check_fat(int FatSz) {
    return memcmp(SECTOR_HEAD(img_buffer, 1), SECTOR_HEAD(img_buffer, 1 + FatSz), FatSz * 512);
}

static int get_next_cluster(int FatNumber, int cur_cluster, BPB_type* bpb_ptr) {
    int FATOffset = cur_cluster + (cur_cluster / 2);
    int ThisFATSecNum = bpb_ptr->RsvdSecCnt + (FATOffset / bpb_ptr->BytsPerSec);
    int ThisFATEntOffset = FATOffset % bpb_ptr->BytsPerSec;
    int SectorNumber = (FatNumber * bpb_ptr->FATSz16) + ThisFATSecNum;
    int TotalOffset = SectorNumber * bpb_ptr->BytsPerSec + ThisFATEntOffset;
    int FAT12ClusEntryVal = *((uint16_t *) &img_buffer[TotalOffset]);
    if (cur_cluster & 0x0001) {
        FAT12ClusEntryVal = FAT12ClusEntryVal >> 4; /* Cluster number is ODD */
    }
    else {
        FAT12ClusEntryVal = FAT12ClusEntryVal & 0x0FFF; /* Cluster number is EVEN */
    }
    return FAT12ClusEntryVal;
}

static void print_cluster_link(int FatNumber, int cluster, BPB_type* bpb_ptr) {
    printf("Cluster link: ");
    while (1)
    {
        int next_cluster = get_next_cluster(FatNumber, cluster, bpb_ptr);
        if (next_cluster == CLUSTER_BAD) {
            printf("Bad Cluster %d\n", cluster);
            assert(0);
        }
        else if (next_cluster == CLUSTER_FREE) {
            printf("Free Cluster %d\n", cluster);
            break;
        }
        else if (CLUSTER_RSRVD <= next_cluster && next_cluster <= CLUSTER_FINAL) {
            printf("%d\n", cluster);
            break;
        }
        else {
            printf("%d ", cluster);
            cluster = next_cluster;
        }
    }
}

static void print_dir_entry(DIR_type* dir_ptr) {
    printf("DIR/FILE:\n");
    printf("Name:         "); for (int i = 0; i < 11; ++ i) putchar(dir_ptr->Name[i]); printf("\n");
    printf("Attr:         0x%02x\n", dir_ptr->Attr);
    printf("NTRes:        %u\n", dir_ptr->NTRes);
    printf("CrtTimeTenth: %u\n", dir_ptr->CrtTimeTenth);
    printf("CrtTime:      %u\n", dir_ptr->CrtTime);
    printf("CrtDate:      %u\n", dir_ptr->CrtDate);
    printf("LstAccDate:   %u\n", dir_ptr->LstAccDate);
    printf("FstClusHI:    %u\n", dir_ptr->FstClusHI);
    printf("WrtTime:      %u\n", dir_ptr->WrtTime);
    printf("WrtDate:      %u\n", dir_ptr->WrtDate);
    printf("FstClusLO:    %u\n", dir_ptr->FstClusLO);
    printf("FileSize:     %u\n", dir_ptr->FileSize);
}

static inline void sector_copy(int dst, int src, int n) {
    memcpy(SECTOR_HEAD(img_buffer, dst), SECTOR_HEAD(img_buffer, src), n * 512);
}

static void print_sector_as_dir(int sector, BPB_type *bpb_ptr) {
    DIR_type* sector_ptr = (DIR_type* )SECTOR_HEAD(img_buffer, sector);
    for (int i = 0; i < bpb_ptr->BytsPerSec / 32; ++ i) {
        DIR_type* dir_ptr = &sector_ptr[i];
        if (dir_ptr->Name[0]) {
            print_dir_entry(dir_ptr);
        }
    }
}

static void find_all_files(BPB_type* bpb_ptr) {
    static uint8_t vis[3000] = {0};
    static int possible_files[] = {526, 343, 14, 11, 7};
    int count = 0;
    for (int i = 0; i <= 5; ++ i) {
        if (vis[possible_files[i]]) continue;
        if (get_next_cluster(0, i, bpb_ptr)) {
            printf("Cluster link: ");
            int cluster = possible_files[i];
            count = 0;
            while (1)
            {
                count ++;
                int next_cluster = get_next_cluster(0, cluster, bpb_ptr);
                if (next_cluster == CLUSTER_BAD) {
                    printf("Bad Cluster %d\n", cluster);
                    assert(0);
                }
                else if (next_cluster == CLUSTER_FREE) {
                    printf("Free Cluster %d\n", cluster);
                    break;
                }
                else if (CLUSTER_RSRVD <= next_cluster && next_cluster <= CLUSTER_FINAL) {
                    printf("%d\n", cluster);
                    vis[cluster] = 1;
                    break;
                }
                else {
                    printf("%d ", cluster);
                    cluster = next_cluster;
                    vis[cluster] = 1;
                }
            }
            printf("Sector Count: %d\n", count);
        }
    }
}

int main() {
    read_img("broken_disk-edit2.img");
    BPB_type* bpb_ptr;
    bpb_ptr = (BPB_type*)img_buffer;
    print_bpb(bpb_ptr);
    print_cluster_link(0, 5, bpb_ptr);
    int RootDirSectors = ((bpb_ptr->RootEntCnt * 32) + (bpb_ptr->BytsPerSec - 1)) / bpb_ptr->BytsPerSec;
    printf("RootDirSectors: %d\n", RootDirSectors);
    find_all_files(bpb_ptr);
    // int root_dir_sector_num = bpb_ptr->RsvdSecCnt + bpb_ptr->NumFATs * bpb_ptr->FATSz16;
    // print_sector_as_dir(root_dir_sector_num, bpb_ptr);
    free(img_buffer);
    return 0;
}