// fat32.c - FAT32 filesystem driver
// Registers itself via IoRegisterFilesystem after finding a FAT32 partition.
// Supports read-only access; write support is not implemented.
#include <io_driver.h>
#include <hal_api.h>
#include <io_file.h>
#include <mm_virt.h>
#include <ke_debug.h>
#include <stdint.h>
#include <stddef.h>
#include <string.h>

// ---------------------------------------------------------------------------
// On-disk structures
// ---------------------------------------------------------------------------
#pragma pack(push, 1)

typedef struct {
    uint8_t  Jump[3];
    char     OemId[8];
    uint16_t BytesPerSector;
    uint8_t  SectorsPerCluster;
    uint16_t ReservedSectors;
    uint8_t  NumFats;
    uint16_t RootEntryCount;     // 0 for FAT32
    uint16_t TotalSectors16;
    uint8_t  MediaType;
    uint16_t FatSize16;          // 0 for FAT32
    uint16_t SectorsPerTrack;
    uint16_t NumHeads;
    uint32_t HiddenSectors;
    uint32_t TotalSectors32;
    // FAT32 extended BPB
    uint32_t FatSize32;
    uint16_t ExtFlags;
    uint16_t FsVersion;
    uint32_t RootCluster;
    uint16_t FsInfoSector;
    uint16_t BackupBootSector;
    uint8_t  Reserved[12];
    uint8_t  DriveNumber;
    uint8_t  Reserved1;
    uint8_t  BootSignature;
    uint32_t VolumeId;
    char     VolumeLabel[11];
    char     FsType[8];
} BPB;

typedef struct {
    char     Name[8];
    char     Ext[3];
    uint8_t  Attr;
    uint8_t  Reserved;
    uint8_t  CrtTimeTenth;
    uint16_t CrtTime;
    uint16_t CrtDate;
    uint16_t LastAccDate;
    uint16_t FirstClusterHi;
    uint16_t WrtTime;
    uint16_t WrtDate;
    uint16_t FirstClusterLo;
    uint32_t FileSize;
} DIRENT;

typedef struct {
    uint8_t  Order;
    uint16_t Name1[5];
    uint8_t  Attr;       // 0x0F
    uint8_t  Type;
    uint8_t  Checksum;
    uint16_t Name2[6];
    uint16_t FirstCluster;
    uint16_t Name3[2];
} LFN_ENTRY;

// GPT Structures
typedef struct {
    uint32_t Data1;
    uint16_t Data2;
    uint16_t Data3;
    uint8_t  Data4[8];
} EFI_GUID;

typedef struct {
    char     Signature[8]; // "EFI PART"
    uint32_t Revision;
    uint32_t HeaderSize;
    uint32_t HeaderCRC32;
    uint32_t Reserved;
    uint64_t MyLBA;
    uint64_t AlternateLBA;
    uint64_t FirstUsableLBA;
    uint64_t LastUsableLBA;
    EFI_GUID DiskGUID;
    uint64_t PartitionEntryLBA;
    uint32_t NumberOfPartitionEntries;
    uint32_t SizeOfPartitionEntry;
    uint32_t PartitionEntryArrayCRC32;
} GPT_HEADER;

typedef struct {
    EFI_GUID PartitionTypeGUID;
    EFI_GUID UniquePartitionGUID;
    uint64_t StartingLBA;
    uint64_t EndingLBA;
    uint64_t Attributes;
    uint16_t PartitionName[36];
} GPT_ENTRY;

#pragma pack(pop)

#define ATTR_DIRECTORY  0x10
#define ATTR_LFN        0x0F
#define FAT32_EOC       0x0FFFFFF8

// ---------------------------------------------------------------------------
// Driver state
// ---------------------------------------------------------------------------
typedef struct {
    uint8_t   DriveIndex;
    uint32_t  PartitionLba;   // LBA of the partition's first sector
    uint32_t  BytesPerSector;
    uint32_t  SectorsPerCluster;
    uint32_t  BytesPerCluster;
    uint32_t  ReservedSectors;
    uint32_t  NumFats;
    uint32_t  FatSize;        // in sectors
    uint32_t  RootCluster;
    uint32_t  DataStart;      // LBA of cluster 2
    const HAL_API* Hal;
} FAT32_VOL;

typedef struct {
    uint32_t  StartCluster;
    uint32_t  FileSize;
    FAT32_VOL* Vol;
} FAT32_FILE;

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------
static KSTATUS ReadSectors(FAT32_VOL* v, uint64_t lba,
                           uint32_t count, void* buf) {    
    if (!v || !v->Hal || !v->Hal->Disk.ReadSectors) {
        return STATUS_INVALID_PARAMETER;
    }
    return v->Hal->Disk.ReadSectors(v->DriveIndex, lba, count, buf);
}

static uint64_t ClusterToLba(FAT32_VOL* v, uint32_t cluster) {
    return v->DataStart + (uint64_t)(cluster - 2) * v->SectorsPerCluster;
}

// Read the FAT entry for a given cluster. Returns 0 on error.
static uint32_t FatEntry(FAT32_VOL* v, uint32_t cluster) {
    uint32_t fatOffset = cluster * 4;
    uint32_t sector    = v->PartitionLba + v->ReservedSectors
                       + fatOffset / v->BytesPerSector;
    uint32_t offInSec  = fatOffset % v->BytesPerSector;

    uint8_t* sec = (uint8_t *)MmAllocatePool(v->BytesPerSector);
    if (!sec) return 0;
    if (!KSUCCESS(ReadSectors(v, sector, 1, sec))) {
        DbgPrint("fat32: FatEntry: ReadSectors failed!\n");
    }
    uint32_t val = *(uint32_t *)(sec + offInSec) & 0x0FFFFFFF;
    MmFreePool(sec);
    return val;
}

// ---------------------------------------------------------------------------
// 8.3 name matching
// ---------------------------------------------------------------------------
static int MatchName83(const char* dirName, const char* dirExt,
                       const char* want) {
    char buf[13];
    int  n = 0;
    for (int i = 0; i < 8 && dirName[i] != ' '; i++) {
        buf[n++] = (char)(dirName[i] >= 'A' && dirName[i] <= 'Z'
                          ? dirName[i] + 32 : dirName[i]);
    }
    if (dirExt[0] != ' ') {
        buf[n++] = '.';
        for (int i = 0; i < 3 && dirExt[i] != ' '; i++) {
            buf[n++] = (char)(dirExt[i] >= 'A' && dirExt[i] <= 'Z'
                              ? dirExt[i] + 32 : dirExt[i]);
        }
    }
    buf[n] = '\0';

    const char* p = buf;
    while (*want && *p) {
        char a = *p++,  b = *want++;
        if (a >= 'A' && a <= 'Z') a += 32;
        if (b >= 'A' && b <= 'Z') b += 32;
        if (a != b) return 0;
    }
    return* want == '\0' && *p == '\0';
}

// ---------------------------------------------------------------------------
// LFN name assembly
// ---------------------------------------------------------------------------
#define MAX_LFN  256

static void LfnAppend(char* lfn, const LFN_ENTRY* e) {
    int seq = (e->Order & 0x3F) - 1;
    if (seq < 0 || seq >= (MAX_LFN / 13)) return;

    int pos = seq * 13;
    uint16_t chars[13];
    int idx = 0;
    
    for (int i = 0; i < 5; i++) chars[idx++] = e->Name1[i];
    for (int i = 0; i < 6; i++) chars[idx++] = e->Name2[i];
    for (int i = 0; i < 2; i++) chars[idx++] = e->Name3[i];

    for (int i = 0; i < 13 && pos < MAX_LFN - 1; i++, pos++) {
        uint16_t c = chars[i];
        if (c == 0x0000 || c == 0xFFFF) { lfn[pos] = '\0'; break; }
        lfn[pos] = (char)(c & 0x7F);
    }
}

static int LfnMatch(const char* lfn, const char* want) {
    while (*lfn && *want) {
        char a = *lfn++, b = *want++;
        if (a >= 'A' && a <= 'Z') a += 32;
        if (b >= 'A' && b <= 'Z') b += 32;
        if (a != b) return 0;
    }
    return* lfn == '\0' && *want == '\0';
}

// ---------------------------------------------------------------------------
// Directory search
// ---------------------------------------------------------------------------
static int SearchDir(FAT32_VOL* v, uint32_t dirCluster,
                     const char* name,
                     uint32_t* outCluster, uint32_t* outSize,
                     uint8_t* outAttr) {
    uint8_t* clusterBuf = (uint8_t *)MmAllocatePool(v->BytesPerCluster);
    if (!clusterBuf) return 0;

    char lfn[MAX_LFN];
    int  hasLfn = 0;
    memset(lfn, 0, sizeof(lfn));

    uint32_t cluster = dirCluster;
    while (cluster < FAT32_EOC) {
        if (!KSUCCESS(ReadSectors(v, ClusterToLba(v, cluster),
                                  v->SectorsPerCluster, clusterBuf))) {
            break;
        }

        uint32_t entries = v->BytesPerCluster / sizeof(DIRENT);
        DIRENT* dir     = (DIRENT *)clusterBuf;

        for (uint32_t i = 0; i < entries; i++) {
            if (dir[i].Name[0] == 0x00) goto done;   // end of directory
            if ((uint8_t)dir[i].Name[0] == 0xE5) {   // deleted
                hasLfn = 0; memset(lfn, 0, sizeof(lfn));
                continue;
            }
            if (dir[i].Attr == ATTR_LFN) {
                LfnAppend(lfn, (LFN_ENTRY *)&dir[i]);
                hasLfn = 1;
                continue;
            }

            int match = 0;
            if (hasLfn && lfn[0]) {
                match = LfnMatch(lfn, name);
            }
            if (!match) {
                match = MatchName83(dir[i].Name, dir[i].Ext, name);
            }

            hasLfn = 0; memset(lfn, 0, sizeof(lfn));

            if (match) {
                *outCluster = ((uint32_t)dir[i].FirstClusterHi << 16) | dir[i].FirstClusterLo;
                *outSize    = dir[i].FileSize;
                *outAttr    = dir[i].Attr;
                MmFreePool(clusterBuf);
                return 1;
            }
        }
        cluster = FatEntry(v, cluster);
    }
done:
    MmFreePool(clusterBuf);
    return 0;
}

// ---------------------------------------------------------------------------
// Path traversal
// ---------------------------------------------------------------------------
static KSTATUS Traverse(FAT32_VOL* v, const char* path,
                        uint32_t* outCluster, uint32_t* outSize) {
    uint32_t cluster = v->RootCluster;
    uint32_t size    = 0;
    uint8_t  attr    = ATTR_DIRECTORY;

    char component[256];
    const char* p = path;

    while (*p) {
        while (*p == '\\' || *p == '/') p++;
        if (!*p) break;

        int len = 0;
        while (*p && *p != '\\' && *p != '/' && len < 255)
            component[len++] = *p++;
        component[len] = '\0';
        if (!len) break;

        if (!(attr & ATTR_DIRECTORY)) return STATUS_NOT_FOUND;

        uint32_t childCluster, childSize;
        uint8_t  childAttr;
        if (!SearchDir(v, cluster, component, &childCluster, &childSize, &childAttr))
            return STATUS_NOT_FOUND;

        cluster = childCluster;
        size    = childSize;
        attr    = childAttr;
    }

    *outCluster = cluster;
    *outSize    = size;
    return STATUS_SUCCESS;
}

// ---------------------------------------------------------------------------
// IO_FILESYSTEM vtable implementations
// ---------------------------------------------------------------------------
static KSTATUS Fat32Open(IO_FILESYSTEM* fs, const char* path, void **fileCtx) {
    FAT32_VOL* v = (FAT32_VOL *)fs->DriverCtx;
    uint32_t cluster, size;
    KSTATUS s = Traverse(v, path, &cluster, &size);
    if (!KSUCCESS(s)) return s;

    FAT32_FILE* f = (FAT32_FILE *)MmAllocatePool(sizeof(FAT32_FILE));
    if (!f) return STATUS_NO_MEMORY;
    f->StartCluster = cluster;
    f->FileSize     = size;
    f->Vol          = v;

    *fileCtx = f;
    return STATUS_SUCCESS;
}

static KSTATUS Fat32Read(IO_FILESYSTEM* fs, void* fileCtx,
                         uint64_t offset, void* buf,
                         size_t size, size_t* bytesRead) {
    (void)fs;
    FAT32_FILE* f = (FAT32_FILE *)fileCtx;
    FAT32_VOL* v = f->Vol;

    if (offset >= f->FileSize) { *bytesRead = 0; return STATUS_SUCCESS; }
    if (offset + size > f->FileSize) size = (size_t)(f->FileSize - offset);

    uint8_t* clusterBuf = (uint8_t *)MmAllocatePool(v->BytesPerCluster);
    if (!clusterBuf) return STATUS_NO_MEMORY;

    size_t  totalRead = 0;
    uint8_t* out      = (uint8_t *)buf;

    uint32_t skipClusters = (uint32_t)(offset / v->BytesPerCluster);
    uint32_t offsetInCluster = (uint32_t)(offset % v->BytesPerCluster);
    uint32_t cluster = f->StartCluster;

    for (uint32_t i = 0; i < skipClusters && cluster < FAT32_EOC; i++)
        cluster = FatEntry(v, cluster);

    while (size > 0 && cluster < FAT32_EOC) {
        if (!KSUCCESS(ReadSectors(v, ClusterToLba(v, cluster),
            v->SectorsPerCluster, clusterBuf))) {
            break;
        }

        size_t avail = v->BytesPerCluster - offsetInCluster;
        size_t copy  = avail < size ? avail : size;

        memcpy(out, clusterBuf + offsetInCluster, copy);
        out             += copy;
        totalRead       += copy;
        size            -= copy;
        offsetInCluster = 0;

        cluster = FatEntry(v, cluster);
    }

    MmFreePool(clusterBuf);
    *bytesRead = totalRead;
    return STATUS_SUCCESS;
}

static uint64_t Fat32GetSize(IO_FILESYSTEM* fs, void* fileCtx) {
    (void)fs;
    return ((FAT32_FILE *)fileCtx)->FileSize;
}

static void Fat32Close(IO_FILESYSTEM* fs, void* fileCtx) {
    (void)fs;
    MmFreePool(fileCtx);
}

// ---------------------------------------------------------------------------
// Partition scan - Supports MBR and GPT
// ---------------------------------------------------------------------------
static const EFI_GUID GUID_EFI_SYSTEM = {
    0xC12A7328, 0xF81F, 0x11D2, {0xBA, 0x4B, 0x00, 0xA0, 0xC9, 0x3E, 0xC9, 0x3B}
};
static const EFI_GUID GUID_MS_BASIC_DATA = {
    0xEBD0A0A2, 0xB9E5, 0x4433, {0x87, 0xC0, 0x68, 0xB6, 0xB7, 0x26, 0x99, 0xC7}
};
static const EFI_GUID GUID_LINUX_BASIC_DATA = {
    0x0FC63DAF, 0x8483, 0x4772, {0x8E, 0x79, 0x3D, 0x69, 0xD8, 0x47, 0x7D, 0xE4}
};

// Validates whether the VBR actually contains a FAT32 BPB
static int IsFat32Vbr(FAT32_VOL* v, uint32_t lba) {
    uint8_t sector[512];
    if (!KSUCCESS(ReadSectors(v, lba, 1, sector))) return 0;
    
    if (sector[510] != 0x55 || sector[511] != 0xAA) return 0;
    
    BPB* bpb = (BPB *)sector;
    if (bpb->FatSize16 != 0 || bpb->RootEntryCount != 0) return 0;
    if (bpb->FatSize32 == 0) return 0;
    
    return 1;
}

static uint32_t FindFat32Partition(FAT32_VOL* v) {
    uint8_t mbr[512];

    KSTATUS s = ReadSectors(v, 0, 1, mbr);
    if (!KSUCCESS(s)) {
        DbgPrint("FAT32: [DEBUG] ReadSectors for MBR failed.\n");
        return 0;
    }
    
    if (mbr[510] != 0x55 || mbr[511] != 0xAA) {
        DbgPrint("FAT32: [DEBUG] Invalid signature at LBA 0.\n");
        return 0;
    }

    // Check for GPT Protective MBR (Type 0xEE)
    int isGpt = 0;
    for (int i = 0; i < 4; i++) {
        if (mbr[0x1BE + i * 16 + 4] == 0xEE) {
            isGpt = 1;
            break;
        }
    }

    if (isGpt) {
        DbgPrint("FAT32: [DEBUG] Protective MBR found. Parsing GPT...\n");
        uint8_t gptSector[512];
        if (!KSUCCESS(ReadSectors(v, 1, 1, gptSector))) return 0;

        GPT_HEADER* gpt = (GPT_HEADER *)gptSector;
        if (memcmp(gpt->Signature, "EFI PART", 8) != 0) {
            DbgPrint("FAT32: [DEBUG] Invalid GPT signature.\n");
            return 0;
        }

        uint64_t entryLBA = gpt->PartitionEntryLBA;
        uint32_t numEntries = gpt->NumberOfPartitionEntries;
        uint32_t entrySize = gpt->SizeOfPartitionEntry;

        if (entrySize < sizeof(GPT_ENTRY)) return 0;

        uint32_t entriesPerSector = 512 / entrySize;
        uint8_t entrySector[512];

        for (uint32_t i = 0; i < numEntries; i++) {
            if (i % entriesPerSector == 0) {
                if (!KSUCCESS(ReadSectors(v, entryLBA + (i / entriesPerSector), 1, entrySector))) {
                    return 0;
                }
            }

            GPT_ENTRY* entry = (GPT_ENTRY *)(entrySector + (i % entriesPerSector) * entrySize);
            
            static const EFI_GUID emptyGuid = {0};
            if (memcmp(&entry->PartitionTypeGUID, &emptyGuid, sizeof(EFI_GUID)) == 0) continue;

            // Microsoft Basic Data maps to FAT, NTFS, and exFAT. VBR validation checks true FAT32 nature.
            if (memcmp(&entry->PartitionTypeGUID, &GUID_EFI_SYSTEM, sizeof(EFI_GUID)) == 0 ||
                memcmp(&entry->PartitionTypeGUID, &GUID_MS_BASIC_DATA, sizeof(EFI_GUID)) == 0 ||
                memcmp(&entry->PartitionTypeGUID, &GUID_LINUX_BASIC_DATA, sizeof(EFI_GUID)) == 0) {
                
                uint32_t lba = (uint32_t)entry->StartingLBA;
                if (IsFat32Vbr(v, lba)) {
                    DbgPrint("FAT32: [DEBUG] Found FAT32 partition via GPT at LBA %u\n", lba);
                    return lba;
                }
            }
        }
        DbgPrint("FAT32: [DEBUG] No FAT32 partition found in GPT.\n");
        return 0;
    }

    // Fallback: MBR Parsing
    for (int i = 0; i < 4; i++) {
        uint8_t* entry = mbr + 0x1BE + i * 16;
        uint8_t  type  = entry[4];
        
        if (type == 0x0B || type == 0x0C) {
            uint32_t lba = *(uint32_t *)(entry + 8);
            if (IsFat32Vbr(v, lba)) {
                DbgPrint("FAT32: [DEBUG] Found FAT32 partition via MBR at LBA %u\n", lba);
                return lba;
            }
        }
    }
    
    DbgPrint("FAT32: [DEBUG] No FAT32 partition found in MBR.\n");
    return 0; 
}

// ---------------------------------------------------------------------------
// DriverEntry
// ---------------------------------------------------------------------------
static FAT32_VOL    gVol;
static IO_FILESYSTEM gFs;

KSTATUS DriverEntry(const HAL_API* hal) {
    DbgPrint("FAT32: Init\n");
    if (!hal) return STATUS_INVALID_PARAMETER;

    gVol.Hal = hal;
    gVol.DriveIndex = 0;
    
    uint32_t partLba = FindFat32Partition(&gVol); 
    if (partLba == 0) {
        DbgPrint("FAT32: [DEBUG] Could not locate a valid FAT32 volume.\n");
        return STATUS_NOT_SUPPORTED;
    }

    gVol.PartitionLba = partLba;

    uint8_t sector[512];
    if (!KSUCCESS(ReadSectors(&gVol, partLba, 1, sector))) {
        DbgPrint("FAT32: [DEBUG] ReadSectors FAILED!\n");
        return STATUS_DEVICE_ERROR;
    }

    BPB* bpb = (BPB *)sector;

    // IsFat32Vbr guarantees these pass, but retained for absolute sanity checking.
    if (bpb->FatSize16 != 0 || bpb->RootEntryCount != 0) {
        DbgPrint("FAT32: [DEBUG] Validation failed - Not a FAT32 volume.\n");
        return STATUS_NOT_SUPPORTED;
    }

    gVol.BytesPerSector    = bpb->BytesPerSector;
    gVol.SectorsPerCluster = bpb->SectorsPerCluster;
    gVol.BytesPerCluster   = bpb->BytesPerSector * bpb->SectorsPerCluster;
    gVol.ReservedSectors   = bpb->ReservedSectors;
    gVol.NumFats           = bpb->NumFats;
    gVol.FatSize           = bpb->FatSize32;
    gVol.RootCluster       = bpb->RootCluster;
    gVol.DataStart         = partLba
                           + bpb->ReservedSectors
                           + bpb->NumFats * bpb->FatSize32;

    gFs.Open      = Fat32Open;
    gFs.Read      = Fat32Read;
    gFs.GetSize   = Fat32GetSize;
    gFs.Close     = Fat32Close;
    gFs.DriverCtx = &gVol;

    // Populate IO_VOLUME_PARAMS
    IO_VOLUME_PARAMS volParams;
    volParams.DiskIndex = gVol.DriveIndex;
    volParams.PartitionLba = partLba;
    // FAT32 sizes are stored in TotalSectors32 if TotalSectors16 is 0
    volParams.PartitionSize = (bpb->TotalSectors16 != 0) ? bpb->TotalSectors16 : bpb->TotalSectors32;

    // Pass the new volParams structure
    KSTATUS status = IoRegisterFilesystem(gVol.DriveIndex, &gFs, &volParams);
    DbgPrint("FAT32: [DEBUG] pIoRegisterFilesystem returned: %d\n", status);

    DbgPrint("FAT32: Init complete.\n");
    return STATUS_SUCCESS;
}