#include <fs.h>
#include <serial.h>
#include <memory/heap.h>
#include <string_utils.h>

extern EFI_SYSTEM_TABLE *ST;

#pragma pack(push, 1)

struct BPB {
    UINT8  jmpBoot[3];
    UINT8  OEMName[8];
    UINT16 BytsPerSec;
    UINT8  SecPerClus;
    UINT16 RsvdSecCnt;
    UINT8  NumFATs;
    UINT16 RootEntCnt;
    UINT16 TotSec16;
    UINT8  Media;
    UINT16 FATSz16;
    UINT16 SecPerTrk;
    UINT16 NumHeads;
    UINT32 HiddSec;
    UINT32 TotSec32;
    UINT32 FATSz32;
    UINT16 ExtFlags;
    UINT16 FSVer;
    UINT32 RootClus;
    UINT16 FSInfo;
    UINT16 BkBootSec;
    UINT8  Reserved[12];
    UINT8  DrvNum;
    UINT8  Reserved1;
    UINT8  BootSig;
    UINT32 VolID;
    UINT8  VolLab[11];
    UINT8  FilSysType[8];
};

struct GPTHeader {
    UINT64  Signature;
    UINT32  Revision;
    UINT32  HeaderSize;
    UINT32  HeaderCRC32;
    UINT32  Reserved;
    UINT64  MyLBA;
    UINT64  AlternateLBA;
    UINT64  FirstUsableLBA;
    UINT64  LastUsableLBA;
    EFI_GUID DiskGUID;
    UINT64  PartitionEntryLBA;
    UINT32  NumPartitionEntries;
    UINT32  SizeOfPartitionEntry;
    UINT32  PartitionEntryArrayCRC32;
};

struct GPTEntry {
    EFI_GUID PartitionTypeGUID;
    EFI_GUID UniquePartitionGUID;
    UINT64   StartingLBA;
    UINT64   EndingLBA;
    UINT64   Attributes;
    CHAR16   PartitionName[36];
};

struct MBRPart {
    UINT8   Status;
    UINT8   CHSStart[3];
    UINT8   Type;
    UINT8   CHSEnd[3];
    UINT32  LBABegin;
    UINT32  NumSectors;
};

#pragma pack(pop)

#define GPT_SIGNATURE_VAL 0x5452415020494645ULL
#define MBR_SIGNATURE     0xAA55

#define MBR_TYPE_FAT12    0x01
#define MBR_TYPE_FAT16    0x04
#define MBR_TYPE_FAT16B   0x06
#define MBR_TYPE_FAT32    0x0B
#define MBR_TYPE_FAT32LBA 0x0C
#define MBR_TYPE_FAT16LBA 0x0E
#define MBR_TYPE_ESP      0xEF

// ---- GUID helpers ----

static BOOLEAN guid_is_zero(EFI_GUID *g) {
    return g->Data1 == 0 && g->Data2 == 0 && g->Data3 == 0
        && g->Data4[0] == 0 && g->Data4[1] == 0
        && g->Data4[2] == 0 && g->Data4[3] == 0
        && g->Data4[4] == 0 && g->Data4[5] == 0
        && g->Data4[6] == 0 && g->Data4[7] == 0;
}

static BOOLEAN guid_eq(EFI_GUID *a, EFI_GUID *b) {
    return a->Data1 == b->Data1 && a->Data2 == b->Data2
        && a->Data3 == b->Data3
        && a->Data4[0] == b->Data4[0] && a->Data4[1] == b->Data4[1]
        && a->Data4[2] == b->Data4[2] && a->Data4[3] == b->Data4[3]
        && a->Data4[4] == b->Data4[4] && a->Data4[5] == b->Data4[5]
        && a->Data4[6] == b->Data4[6] && a->Data4[7] == b->Data4[7];
}

// ---- Block I/O ----

struct block_dev {
    EFI_BLOCK_IO_PROTOCOL *Bio;
    UINT32 BlockSize;
};

static EFI_STATUS blk_init(struct block_dev *dev) {
    EFI_GUID g = EFI_BLOCK_IO_PROTOCOL_GUID;
    EFI_HANDLE *Handles = NULL;
    UINTN NoHandles = 0;
    EFI_STATUS st = uefi_call_wrapper(ST->BootServices->LocateHandleBuffer, 5,
        ByProtocol, &g, NULL, &NoHandles, &Handles);
    if (EFI_ERROR(st) || NoHandles == 0) return EFI_NOT_FOUND;

    EFI_BLOCK_IO_PROTOCOL *Bio = NULL;
    for (UINTN i = 0; i < NoHandles; i++) {
        EFI_BLOCK_IO_PROTOCOL *b = NULL;
        st = uefi_call_wrapper(ST->BootServices->HandleProtocol, 3,
            Handles[i], &g, (void**)&b);
        if (EFI_ERROR(st)) continue;
        if (b->Media->MediaPresent) { Bio = b; break; }
        if (!Bio) Bio = b;
    }

    uefi_call_wrapper(ST->BootServices->FreePool, 1, Handles);
    if (!Bio) return EFI_NOT_FOUND;

    dev->Bio = Bio;
    dev->BlockSize = Bio->Media->BlockSize;
    return EFI_SUCCESS;
}

static EFI_STATUS blk_read(struct block_dev *dev, UINT64 LBA,
                            UINTN Sectors, void *Buf) {
    return uefi_call_wrapper(dev->Bio->ReadBlocks, 5,
        dev->Bio, dev->Bio->Media->MediaId, LBA,
        Sectors * dev->BlockSize, Buf);
}

static EFI_STATUS blk_write(struct block_dev *dev, UINT64 LBA,
                             UINTN Sectors, const void *Buf) {
    return uefi_call_wrapper(dev->Bio->WriteBlocks, 5,
        dev->Bio, dev->Bio->Media->MediaId, LBA,
        Sectors * dev->BlockSize, (void*)Buf);
}

// ---- Partition detection ----

static EFI_STATUS find_gpt_partition(struct block_dev *dev, UINT64 *StartLBA) {
    UINT8 *Buf = (UINT8*)kmalloc(dev->BlockSize);
    if (!Buf) return EFI_OUT_OF_RESOURCES;

    EFI_STATUS st = blk_read(dev, 1, 1, Buf);
    if (EFI_ERROR(st)) { kfree(Buf); return st; }

    struct GPTHeader *Hdr = (struct GPTHeader*)Buf;
    if (Hdr->Signature != GPT_SIGNATURE_VAL) { kfree(Buf); return EFI_NOT_FOUND; }

    // Read partition entries
    UINTN EntrySz = Hdr->SizeOfPartitionEntry;
    UINTN NumEnt = Hdr->NumPartitionEntries;
    UINTN Total = EntrySz * NumEnt;
    UINTN Secs = (Total + dev->BlockSize - 1) / dev->BlockSize;

    UINT8 *EntBuf = (UINT8*)kmalloc(Secs * dev->BlockSize);
    if (!EntBuf) { kfree(Buf); return EFI_OUT_OF_RESOURCES; }

    st = blk_read(dev, Hdr->PartitionEntryLBA, Secs, EntBuf);
    if (EFI_ERROR(st)) { kfree(Buf); kfree(EntBuf); return st; }

    // EFI System Partition GUID
    EFI_GUID esp = { 0xC12A7328, 0xF81F, 0x11D2,
        {0xBA, 0x4B, 0x00, 0xA0, 0xC9, 0x3E, 0xC9, 0x3B} };
    // Basic data partition GUID (FAT)
    EFI_GUID basic = { 0xEBD0A0A2, 0xB9E5, 0x4433,
        {0x87, 0xC0, 0x68, 0xB6, 0xB7, 0x26, 0x99, 0xC7} };

    BOOLEAN found = FALSE;
    for (UINTN i = 0; i < NumEnt; i++) {
        struct GPTEntry *E = (struct GPTEntry*)(EntBuf + i * EntrySz);
        if (guid_is_zero(&E->PartitionTypeGUID)) continue;
        if (guid_eq(&E->PartitionTypeGUID, &esp) ||
            guid_eq(&E->PartitionTypeGUID, &basic)) {
            *StartLBA = E->StartingLBA;
            found = TRUE;
            break;
        }
    }

    if (!found) {
        // Fallback: first non-empty partition
        for (UINTN i = 0; i < NumEnt; i++) {
            struct GPTEntry *E = (struct GPTEntry*)(EntBuf + i * EntrySz);
            if (!guid_is_zero(&E->PartitionTypeGUID)) {
                *StartLBA = E->StartingLBA;
                found = TRUE;
                break;
            }
        }
    }

    kfree(Buf);
    kfree(EntBuf);
    return found ? EFI_SUCCESS : EFI_NOT_FOUND;
}

static EFI_STATUS find_mbr_partition(struct block_dev *dev, UINT64 *StartLBA) {
    UINT8 *Buf = (UINT8*)kmalloc(dev->BlockSize);
    if (!Buf) return EFI_OUT_OF_RESOURCES;

    EFI_STATUS st = blk_read(dev, 0, 1, Buf);
    if (EFI_ERROR(st)) { kfree(Buf); return st; }

    if (*(UINT16*)(Buf + 510) != MBR_SIGNATURE) { kfree(Buf); return EFI_NOT_FOUND; }

    struct MBRPart *Parts = (struct MBRPart*)(Buf + 446);
    // Verify at least one non-zero partition entry exists
    BOOLEAN has_part = FALSE;
    for (SSINT32 i = 0; i < 4; i++) {
        if (Parts[i].Type != 0x00 && Parts[i].Type != 0xEE) {
            has_part = TRUE;
            break;
        }
    }
    if (!has_part) { kfree(Buf); return EFI_NOT_FOUND; }

    for (SSINT32 i = 0; i < 4; i++) {
        UINT8 t = Parts[i].Type;
        if (t == 0x00 || t == 0xEE) continue;
        if (t == MBR_TYPE_FAT12 || t == MBR_TYPE_FAT16 ||
            t == MBR_TYPE_FAT16B || t == MBR_TYPE_FAT32 ||
            t == MBR_TYPE_FAT32LBA || t == MBR_TYPE_FAT16LBA ||
            t == MBR_TYPE_ESP) {
            *StartLBA = Parts[i].LBABegin;
            kfree(Buf);
            return EFI_SUCCESS;
        }
    }

    // Fallback: first non-empty partition
    for (SSINT32 i = 0; i < 4; i++) {
        if (Parts[i].Type != 0x00) {
            *StartLBA = Parts[i].LBABegin;
            kfree(Buf);
            return EFI_SUCCESS;
        }
    }

    kfree(Buf);
    return EFI_NOT_FOUND;
}

// ---- FAT driver ----

struct fat_fs {
    struct block_dev *Dev;
    UINT64  PartLBA;
    UINT16  BytsPerSec;
    UINT8   SecPerClus;
    UINT16  RsvdSecCnt;
    UINT8   NumFATs;
    UINT32  FATSz;
    UINT32  RootClus;
    UINT32  TotClus;
    UINT16  RootEntCnt;
    BOOLEAN IsFAT32;
    UINT32  ClusSize;

    void   *FatBuf;   // cached FAT
    UINT32  FATEntries;
    UINT32 *FAT32;
    UINT16 *FAT16;
};

static UINT64 clus_to_lba(struct fat_fs *fs, UINT32 Clus) {
    UINT32 RootSecs = fs->IsFAT32 ? 0
        : ((fs->RootEntCnt * 32) + fs->BytsPerSec - 1) / fs->BytsPerSec;
    UINT32 FirstDataSec = fs->RsvdSecCnt + fs->NumFATs * fs->FATSz + RootSecs;
    return fs->PartLBA + FirstDataSec + (UINT64)(Clus - 2) * fs->SecPerClus;
}

static UINT32 fat_next(struct fat_fs *fs, UINT32 Clus) {
    if (fs->IsFAT32) {
        if (Clus >= fs->FATEntries) return 0x0FFFFFF8;
        UINT32 v = fs->FAT32[Clus] & 0x0FFFFFFF;
        if (v == 0x00000000) return 0;
        if (v >= 0x0FFFFFF8) return 0x0FFFFFF8;
        return v;
    } else {
        if (Clus >= fs->FATEntries) return 0xFFF8;
        UINT16 v = fs->FAT16[Clus];
        if (v == 0x0000) return 0;
        if (v >= 0xFFF8) return 0x0FFFFFF8;
        return v;
    }
}

static EFI_STATUS fat_init(struct fat_fs *fs, struct block_dev *dev, UINT64 PartLBA) {
    fs->Dev = dev;
    fs->PartLBA = PartLBA;

    UINT8 *Buf = (UINT8*)kmalloc(dev->BlockSize);
    if (!Buf) return EFI_OUT_OF_RESOURCES;

    EFI_STATUS st = blk_read(dev, PartLBA, 1, Buf);
    if (EFI_ERROR(st)) { kfree(Buf); return st; }

    if (*(UINT16*)(Buf + 510) != MBR_SIGNATURE) { kfree(Buf); return EFI_UNSUPPORTED; }

    struct BPB *bpb = (struct BPB*)Buf;
    fs->BytsPerSec = bpb->BytsPerSec;
    fs->SecPerClus = bpb->SecPerClus;
    fs->RsvdSecCnt = bpb->RsvdSecCnt;
    fs->NumFATs = bpb->NumFATs;
    fs->RootEntCnt = bpb->RootEntCnt;

    UINT32 TotSec = bpb->TotSec16 ? bpb->TotSec16 : bpb->TotSec32;
    UINT32 RootSecs = ((fs->RootEntCnt * 32) + fs->BytsPerSec - 1) / fs->BytsPerSec;
    UINT32 FATSz = bpb->FATSz16 ? bpb->FATSz16 : bpb->FATSz32;
    fs->FATSz = FATSz;

    UINT32 FirstDataSec = fs->RsvdSecCnt + fs->NumFATs * FATSz + RootSecs;
    fs->TotClus = (TotSec - FirstDataSec) / fs->SecPerClus;

    if (fs->TotClus < 4085) {
        serial_write("Sylva: FS: FAT12 unsupported\n");
        kfree(Buf);
        return EFI_UNSUPPORTED;
    } else if (fs->TotClus < 65525) {
        fs->IsFAT32 = FALSE;
        fs->FATEntries = FATSz * fs->BytsPerSec / 2;
    } else {
        fs->IsFAT32 = TRUE;
        fs->RootClus = bpb->RootClus;
        fs->FATEntries = FATSz * fs->BytsPerSec / 4;
    }

    fs->ClusSize = fs->BytsPerSec * fs->SecPerClus;

    // Cache FAT
    UINTN FATBytes = FATSz * fs->BytsPerSec;
    fs->FatBuf = kmalloc(FATBytes);
    if (!fs->FatBuf) { kfree(Buf); return EFI_OUT_OF_RESOURCES; }

    st = blk_read(dev, PartLBA + fs->RsvdSecCnt, FATSz, fs->FatBuf);
    if (EFI_ERROR(st)) { kfree(Buf); kfree(fs->FatBuf); return st; }

    if (fs->IsFAT32) {
        fs->FAT32 = (UINT32*)fs->FatBuf;
        fs->FAT16 = NULL;
    } else {
        fs->FAT16 = (UINT16*)fs->FatBuf;
        fs->FAT32 = NULL;
    }

    kfree(Buf);
    return EFI_SUCCESS;
}

// ---- LFN helpers ----

static UINT8 lfn_checksum(const UINT8 *SFN) {
    UINT8 sum = 0;
    for (SSINT32 i = 0; i < 11; i++)
        sum = ((sum & 1) ? 0x80 : 0) + (sum >> 1) + SFN[i];
    return sum;
}

#define LFN_MAX_FRAGS 20
#define LFN_FRAG_SIZE 13

struct lfn_state {
    CHAR16 frags[LFN_MAX_FRAGS][LFN_FRAG_SIZE + 1];
    UINTN  count;
    UINT8  checksum;
};

static void lfn_reset(struct lfn_state *lfn) {
    lfn->count = 0;
    lfn->checksum = 0;
}

static void lfn_add(struct lfn_state *lfn, const UINT8 *E) {
    if (lfn->count >= LFN_MAX_FRAGS) return;
    lfn->checksum = E[13];
    UINTN pos = 0;
    BOOLEAN done = FALSE;
    for (SSINT32 i = 0; i < 5 && pos < LFN_FRAG_SIZE && !done; i++) {
        CHAR16 c = *(const UINT16*)(E + 1 + i * 2);
        if (c == 0x0000 || c == 0xFFFF) { done = TRUE; break; }
        lfn->frags[lfn->count][pos++] = c;
    }
    for (SSINT32 i = 0; i < 6 && pos < LFN_FRAG_SIZE && !done; i++) {
        CHAR16 c = *(const UINT16*)(E + 14 + i * 2);
        if (c == 0x0000 || c == 0xFFFF) { done = TRUE; break; }
        lfn->frags[lfn->count][pos++] = c;
    }
    for (SSINT32 i = 0; i < 2 && pos < LFN_FRAG_SIZE && !done; i++) {
        CHAR16 c = *(const UINT16*)(E + 28 + i * 2);
        if (c == 0x0000 || c == 0xFFFF) { done = TRUE; break; }
        lfn->frags[lfn->count][pos++] = c;
    }
    lfn->frags[lfn->count][pos] = 0;
    lfn->count++;
}

// Fragments arrive in reverse order (seq=N first, seq=1 last).
// Concatenate from last to first to get correct name order.
static void lfn_build(struct lfn_state *lfn, CHAR16 *out, UINTN out_size) {
    UINTN pos = 0;
    for (UINTN i = lfn->count; i > 0; i--) {
        for (UINTN j = 0; lfn->frags[i - 1][j] && pos < out_size - 1; j++)
            out[pos++] = lfn->frags[i - 1][j];
    }
    out[pos] = 0;
}

static void sfn_to_name(const UINT8 *E, CHAR16 *out, UINTN out_size) {
    UINTN pos = 0;
    for (SSINT32 i = 0; i < 8 && pos < out_size - 1; i++)
        if (E[i] != ' ') out[pos++] = E[i];
    UINTN ext_start = pos;
    BOOLEAN has_ext = FALSE;
    for (SSINT32 i = 8; i < 11 && pos < out_size - 1; i++) {
        if (E[i] != ' ') {
            if (!has_ext) { out[pos++] = '.'; has_ext = TRUE; ext_start = pos; }
            out[pos++] = E[i];
        }
    }
    // If extension is empty but we added a dot, remove it
    if (has_ext && pos == ext_start) pos--;
    out[pos] = 0;
}

// ---- Directory reading ----

typedef void (*dir_cb)(void *Ctx, const CHAR16 *Name, UINT8 Attr,
                        UINT32 Size, UINT32 FirstClus, UINTN EntryOff);

// Returns TRUE when end-of-directory reached
static BOOLEAN process_sector(struct fat_fs *fs, UINT8 *Buf, UINTN SectBase,
                                struct lfn_state *lfn,
                                dir_cb Callback, void *Ctx) {
    for (UINTN off = 0; off < fs->BytsPerSec; off += 32) {
        UINT8 *E = Buf + off;

        if (E[0] == 0x00) return TRUE;
        if (E[0] == 0xE5) { lfn_reset(lfn); continue; }

        UINT8 Attr = E[11];

        if ((Attr & 0x3F) == 0x0F) {
            if (E[0] & 0x40) lfn_reset(lfn);
            lfn_add(lfn, E);
            continue;
        }

        // SFN entry
        CHAR16 Name[256];
        BOOLEAN use_lfn = FALSE;
        if (lfn->count > 0 && lfn_checksum(E) == lfn->checksum) {
            lfn_build(lfn, Name, 256);
            use_lfn = TRUE;
        }
        if (!use_lfn) sfn_to_name(E, Name, 256);

        // Skip . and ..
        if (E[0] == 0x2E) { lfn_reset(lfn); continue; }

        UINT32 Size = *(const UINT32*)(E + 28);
        UINT32 FirstClus;
        if (fs->IsFAT32)
            FirstClus = ((UINT32)*(const UINT16*)(E + 20) << 16)
                       | *(const UINT16*)(E + 26);
        else
            FirstClus = *(const UINT16*)(E + 26);

        Callback(Ctx, Name, Attr, Size, FirstClus, SectBase + off);
        lfn_reset(lfn);
    }
    return FALSE;
}

static void read_directory(struct fat_fs *fs, UINT32 Cluster,
                            dir_cb Callback, void *Ctx) {
    UINTN ClusBytes = fs->ClusSize;
    UINT8 *Buf = (UINT8*)kmalloc(ClusBytes);
    if (!Buf) return;

    struct lfn_state lfn;
    lfn_reset(&lfn);

    if (fs->IsFAT32) {
        UINT32 Clus = Cluster;
        while (Clus >= 2 && Clus < 0x0FFFFFF8) {
            UINT64 BaseLBA = clus_to_lba(fs, Clus);
            if (EFI_ERROR(blk_read(fs->Dev, BaseLBA, fs->SecPerClus, Buf)))
                goto done;
            for (UINTN s = 0; s < fs->SecPerClus; s++) {
                if (process_sector(fs, Buf + s * fs->BytsPerSec,
                                    s * fs->BytsPerSec,
                                    &lfn, Callback, Ctx))
                    goto done;
            }
            Clus = fat_next(fs, Clus);
        }
    } else if (Cluster == 0) {
        UINT32 RootSecs = ((fs->RootEntCnt * 32) + fs->BytsPerSec - 1) / fs->BytsPerSec;
        UINT64 RootLBA = fs->PartLBA + fs->RsvdSecCnt + fs->NumFATs * fs->FATSz;
        for (UINTN s = 0; s < RootSecs; s++) {
            if (EFI_ERROR(blk_read(fs->Dev, RootLBA + s, 1, Buf + s * fs->BytsPerSec)))
                goto done;
        }
        for (UINTN s = 0; s < RootSecs; s++) {
            if (process_sector(fs, Buf + s * fs->BytsPerSec,
                                s * fs->BytsPerSec,
                                &lfn, Callback, Ctx))
                goto done;
        }
    } else {
        UINT32 Clus = Cluster;
        while (Clus >= 2 && Clus < 0x0FFFFFF8) {
            UINT64 BaseLBA = clus_to_lba(fs, Clus);
            if (EFI_ERROR(blk_read(fs->Dev, BaseLBA, fs->SecPerClus, Buf)))
                goto done;
            for (UINTN s = 0; s < fs->SecPerClus; s++) {
                if (process_sector(fs, Buf + s * fs->BytsPerSec,
                                    s * fs->BytsPerSec,
                                    &lfn, Callback, Ctx))
                    goto done;
            }
            Clus = fat_next(fs, Clus);
        }
    }

done:
    kfree(Buf);
}

// ---- Public API ----

static struct fat_fs g_fs;
static BOOLEAN g_fs_inited = FALSE;

// Recursive listing support

struct list_ctx {
    struct fat_fs *fs;
    SSINT32 depth;
};

static void name_to_ascii(const CHAR16 *Name, char *Ascii, UINTN ascii_sz) {
    wstr_to_ascii(Ascii, (WString)Name, ascii_sz);
}

static void list_callback(void *ctx, const CHAR16 *Name, UINT8 Attr,
                           UINT32 Size, UINT32 FirstClus, UINTN EntryOff) {
    struct list_ctx *lc = (struct list_ctx*)ctx;
    for (SSINT32 i = 0; i < lc->depth; i++) serial_write("  ");

    serial_write(Attr & 0x10 ? "[DIR]  " : "[FILE] ");

    char ascii[256];
    name_to_ascii(Name, ascii, sizeof(ascii));
    serial_write(ascii);

    if (!(Attr & 0x10)) {
        serial_write(" (");
        serial_write_hex(Size);
        serial_write(" bytes)");
    }
    serial_write("\n");

    if ((Attr & 0x10) && lc->depth < 8) {
        struct list_ctx sub = { lc->fs, lc->depth + 1 };
        read_directory(lc->fs, FirstClus, list_callback, &sub);
    }
}

EFI_STATUS fs_init() {
    serial_write("Sylva: FS: init block...\n");
    static struct block_dev dev;
    EFI_STATUS st = blk_init(&dev);
    if (EFI_ERROR(st)) {
        serial_write("Sylva: FS: no block device!\n");
        return st;
    }

    serial_write("Sylva: FS: block size = ");
    serial_write_hex(dev.BlockSize);
    serial_write("\n");

    UINT64 PartLBA = 0;
    st = find_gpt_partition(&dev, &PartLBA);
    if (EFI_ERROR(st)) {
        serial_write("Sylva: FS: no GPT, trying MBR...\n");
        st = find_mbr_partition(&dev, &PartLBA);
    }
    if (EFI_ERROR(st)) {
        serial_write("Sylva: FS: no partition table, assuming super-floppy...\n");
        PartLBA = 0;
    }

    serial_write("Sylva: FS: partition LBA = ");
    serial_write_hex(PartLBA);
    serial_write("\n");

    st = fat_init(&g_fs, &dev, PartLBA);
    if (EFI_ERROR(st)) {
        serial_write("Sylva: FS: FAT init failed!\n");
        return st;
    }

    serial_write("Sylva: FS: FAT");
    serial_write(g_fs.IsFAT32 ? "32" : "16");
    serial_write(" ready, cluster size = ");
    serial_write_hex(g_fs.ClusSize);
    serial_write("\n");

    g_fs_inited = TRUE;
    return EFI_SUCCESS;
}

void fs_list() {
    if (!g_fs_inited) { serial_write("FS not initialized\n"); return; }
    UINT32 Root = g_fs.IsFAT32 ? g_fs.RootClus : 0;
    struct list_ctx lc = { &g_fs, 0 };
    read_directory(&g_fs, Root, list_callback, &lc);
}

// ---- File reading ----

struct find_ctx {
    const CHAR16 *Target;
    BOOLEAN Found;
    UINT32 Cluster;
    UINT32 Size;
    UINT8  Attr;
};

static BOOLEAN name_match(const CHAR16 *a, const CHAR16 *b) {
    while (*a && *b) {
        CHAR16 ca = *a, cb = *b;
        if (ca >= L'A' && ca <= L'Z') ca += 32;
        if (cb >= L'A' && cb <= L'Z') cb += 32;
        if (ca != cb) return FALSE;
        a++; b++;
    }
    return *a == 0 && *b == 0;
}

static void find_callback(void *ctx, const CHAR16 *Name, UINT8 Attr,
                           UINT32 Size, UINT32 FirstClus, UINTN EntryOff) {
    struct find_ctx *fc = (struct find_ctx*)ctx;
    if (fc->Found) return;
    if (name_match(fc->Target, Name)) {
        fc->Found = TRUE;
        fc->Cluster = FirstClus;
        fc->Size = Size;
        fc->Attr = Attr;
    }
}

EFI_STATUS fs_read(WString Path, void **Buffer, UINTN *Size) {
    if (!g_fs_inited) return EFI_NOT_READY;

    const CHAR16 *p = Path;
    while (*p == L'\\') p++;

    UINT32 CurClus = g_fs.IsFAT32 ? g_fs.RootClus : 0;

    CHAR16 Comp[256];
    while (*p) {
        UINTN ci = 0;
        while (*p && *p != L'\\' && ci < 255) Comp[ci++] = *p++;
        Comp[ci] = 0;
        while (*p == L'\\') p++;

        struct find_ctx fc;
        fc.Target = Comp;
        fc.Found = FALSE;
        read_directory(&g_fs, CurClus, find_callback, &fc);
        if (!fc.Found) return EFI_NOT_FOUND;

        if (*p == 0 && !(fc.Attr & 0x10)) {
            UINT32 FileSz = fc.Size;
            void *Buf = kmalloc(FileSz ? FileSz : 1);
            if (!Buf) return EFI_OUT_OF_RESOURCES;

            if (FileSz > 0) {
                UINT32 Clus = fc.Cluster;
                UINTN Offset = 0;
                UINTN ClusBytes = g_fs.ClusSize;

                while (Clus >= 2 && Clus < 0x0FFFFFF8 && Offset < FileSz) {
                    // Find longest contiguous run starting at Clus
                    UINT32 RunStart = Clus;
                    UINT32 RunEnd = Clus;
                    UINT32 Next = fat_next(&g_fs, Clus);
                    while (Next == RunEnd + 1 && Next >= 2 && Next < 0x0FFFFFF8) {
                        RunEnd = Next;
                        Next = fat_next(&g_fs, RunEnd);
                    }
                    UINTN RunClus = (UINTN)(RunEnd - RunStart + 1);
                    UINT64 LBA = clus_to_lba(&g_fs, RunStart);
                    UINTN Want = RunClus * ClusBytes;
                    if (Want > FileSz - Offset) Want = FileSz - Offset;
                    UINTN NSecs = (Want + g_fs.BytsPerSec - 1) / g_fs.BytsPerSec;
                    EFI_STATUS st = blk_read(g_fs.Dev, LBA, NSecs,
                                              (UINT8*)Buf + Offset);
                    if (EFI_ERROR(st)) { kfree(Buf); return st; }
                    Offset += Want;
                    Clus = Next;
                }
            }

            *Buffer = Buf;
            *Size = FileSz;
            return EFI_SUCCESS;
        }

        if (!(fc.Attr & 0x10)) return EFI_NOT_FOUND;
        CurClus = fc.Cluster;
    }

    return EFI_NOT_FOUND;
}

// ---- FAT allocation ----

static BOOLEAN fat_entry_free(struct fat_fs *fs, UINT32 Clus) {
    if (Clus >= fs->FATEntries) return FALSE;
    if (fs->IsFAT32) return (fs->FAT32[Clus] & 0x0FFFFFFF) == 0;
    return fs->FAT16[Clus] == 0;
}

static void fat_entry_set(struct fat_fs *fs, UINT32 Clus, UINT32 Val) {
    if (Clus >= fs->FATEntries) return;
    if (fs->IsFAT32) fs->FAT32[Clus] = Val;
    else fs->FAT16[Clus] = (UINT16)Val;
}

// Allocate `Count` free clusters and link them into a chain.
// Returns the first cluster of the new chain; the last entry is marked EOF.
static EFI_STATUS fat_alloc_chain(struct fat_fs *fs, UINT32 Count, UINT32 *OutFirst) {
    UINT32 Found = 0;
    UINT32 First = 0, Prev = 0;
    UINT32 EOC = fs->IsFAT32 ? 0x0FFFFFFF : 0xFFFF;

    for (UINT32 i = 2; i < fs->FATEntries && Found < Count; i++) {
        if (!fat_entry_free(fs, i)) continue;
        if (Found == 0) First = i;
        else fat_entry_set(fs, Prev, i);
        Found++;
        Prev = i;
    }
    if (Found < Count) return EFI_OUT_OF_RESOURCES;
    fat_entry_set(fs, Prev, EOC);
    *OutFirst = First;
    return EFI_SUCCESS;
}

// Write the in-memory FAT cache back to all FAT copies on disk.
static EFI_STATUS fat_flush(struct fat_fs *fs) {
    for (UINTN f = 0; f < fs->NumFATs; f++) {
        EFI_STATUS st = blk_write(fs->Dev,
            fs->PartLBA + fs->RsvdSecCnt + f * fs->FATSz,
            fs->FATSz, fs->FatBuf);
        if (EFI_ERROR(st)) return st;
    }
    return EFI_SUCCESS;
}

// ---- File writing ----

struct find_write_ctx {
    const CHAR16 *Target;
    BOOLEAN Found;
    UINT32 Cluster;
    UINT32 Size;
    UINT8  Attr;
    UINTN  DirOffset;
};

static void find_write_callback(void *ctx, const CHAR16 *Name, UINT8 Attr,
                                 UINT32 Size, UINT32 FirstClus, UINTN EntryOff) {
    struct find_write_ctx *fc = (struct find_write_ctx*)ctx;
    if (fc->Found) return;
    if (name_match(fc->Target, Name)) {
        fc->Found = TRUE;
        fc->Cluster = FirstClus;
        fc->Size = Size;
        fc->Attr = Attr;
        fc->DirOffset = EntryOff;
    }
}

// Update a dirent's size + first-cluster fields in-place inside `DirBuf`
// at the given byte offset. For FAT32 we also patch the high cluster word.
static void dirent_update_size(void *DirBuf, UINTN EntryOff,
                                UINT32 NewSize, UINT32 FirstClus,
                                BOOLEAN PatchFirstClus, BOOLEAN IsFAT32) {
    UINT8 *E = (UINT8*)DirBuf + EntryOff;
    *(UINT32*)(E + 28) = NewSize;
    if (PatchFirstClus) {
        if (IsFAT32) {
            *(UINT16*)(E + 20) = (UINT16)(FirstClus >> 16);
            *(UINT16*)(E + 26) = (UINT16)(FirstClus & 0xFFFF);
        } else {
            *(UINT16*)(E + 26) = (UINT16)(FirstClus & 0xFFFF);
        }
    }
}

// Write [Data, Data+Size) into the file at Path, starting at byte Offset.
// Semantics: existing file only (no create); Offset must be <= current size.
// New file size = Offset + Size. Trailing clusters beyond the new size are
// left allocated (lost) — no truncate support.
EFI_STATUS fs_write(WString Path, const void *Data, UINTN Size, UINTN Offset) {
    if (!g_fs_inited) return EFI_NOT_READY;
    if (Data == NULL && Size > 0) return EFI_INVALID_PARAMETER;

    const CHAR16 *p = Path;
    while (*p == L'\\') p++;

    UINT32 CurClus = g_fs.IsFAT32 ? g_fs.RootClus : 0;

    CHAR16 Comp[256];
    BOOLEAN Resolved = FALSE;
    UINT32 FileFirst = 0, FileSize = 0, FileAttr = 0;
    UINTN  FileDirOff = 0;

    while (*p) {
        UINTN ci = 0;
        while (*p && *p != L'\\' && ci < 255) Comp[ci++] = *p++;
        Comp[ci] = 0;
        while (*p == L'\\') p++;

        struct find_write_ctx fc;
        fc.Target = Comp;
        fc.Found = FALSE;
        read_directory(&g_fs, CurClus, find_write_callback, &fc);
        if (!fc.Found) return EFI_NOT_FOUND;

        if (*p == 0) {
            if (fc.Attr & 0x10) return EFI_INVALID_PARAMETER;
            Resolved = TRUE;
            FileFirst = fc.Cluster;
            FileSize = fc.Size;
            FileAttr = fc.Attr;
            FileDirOff = fc.DirOffset;
        } else {
            if (!(fc.Attr & 0x10)) return EFI_NOT_FOUND;
            CurClus = fc.Cluster;
        }
    }
    if (!Resolved) return EFI_NOT_FOUND;
    if (Offset > FileSize) return EFI_INVALID_PARAMETER;

    UINT32 NewSize = (UINT32)(Offset + Size);
    UINT32 OldClusCount = (FileSize + g_fs.ClusSize - 1) / g_fs.ClusSize;
    UINT32 NewClusCount = (NewSize + g_fs.ClusSize - 1) / g_fs.ClusSize;
    UINT32 AddClusCount = NewClusCount - OldClusCount;
    BOOLEAN FileWasEmpty = (FileSize == 0);

    // --- Extend cluster chain if we need more clusters ---
    if (AddClusCount > 0) {
        UINT32 NewFirst;
        EFI_STATUS st = fat_alloc_chain(&g_fs, AddClusCount, &NewFirst);
        if (EFI_ERROR(st)) return st;

        if (FileWasEmpty) {
            FileFirst = NewFirst;
        } else {
            // Walk to the last cluster of the existing chain
            UINT32 Cur = FileFirst;
            UINT32 Next;
            while (TRUE) {
                Next = fat_next(&g_fs, Cur);
                if (Next >= 0x0FFFFFF8) break;
                Cur = Next;
            }
            // Cur is the last cluster; link to new chain
            fat_entry_set(&g_fs, Cur, NewFirst);
        }
    }

    // --- Write data clusters ---
    if (Size > 0) {
        // Walk to the cluster containing byte Offset
        UINT32 Clus = FileFirst;
        UINTN SkipBytes = Offset;
        while (SkipBytes >= g_fs.ClusSize) {
            Clus = fat_next(&g_fs, Clus);
            SkipBytes -= g_fs.ClusSize;
        }

        void *ClusBuf = kmalloc(g_fs.ClusSize);
        if (!ClusBuf) return EFI_OUT_OF_RESOURCES;

        const UINT8 *Src = (const UINT8*)Data;
        UINTN Remaining = Size;
        UINTN InClusOff = SkipBytes;

        while (Remaining > 0) {
            UINT64 LBA = clus_to_lba(&g_fs, Clus);
            UINTN SpaceInClus = g_fs.ClusSize - InClusOff;
            UINTN Chunk = Remaining < SpaceInClus ? Remaining : SpaceInClus;

            if (InClusOff > 0 || Chunk < g_fs.ClusSize) {
                EFI_STATUS rst = blk_read(g_fs.Dev, LBA, g_fs.SecPerClus, ClusBuf);
                if (EFI_ERROR(rst)) { kfree(ClusBuf); return rst; }
                mem_copy((UINT8*)ClusBuf + InClusOff, Src, Chunk);
                EFI_STATUS wst = blk_write(g_fs.Dev, LBA, g_fs.SecPerClus, ClusBuf);
                if (EFI_ERROR(wst)) { kfree(ClusBuf); return wst; }
            } else {
                EFI_STATUS wst = blk_write(g_fs.Dev, LBA, g_fs.SecPerClus, Src);
                if (EFI_ERROR(wst)) { kfree(ClusBuf); return wst; }
            }

            Src += Chunk;
            Remaining -= Chunk;
            InClusOff = 0;
            if (Remaining > 0) {
                Clus = fat_next(&g_fs, Clus);
            }
        }

        kfree(ClusBuf);
    }

    // --- Update dirent (size, and first cluster if file was empty) ---
    {
        void *DirBuf;
        UINT64 DirLBA;
        UINTN  DirNSecs;
        UINTN  DirBufBytes;

        if (g_fs.IsFAT32) {
            DirLBA = clus_to_lba(&g_fs, CurClus);
            DirNSecs = g_fs.SecPerClus;
            DirBufBytes = g_fs.ClusSize;
        } else if (CurClus == 0) {
            UINT32 RootSecs = ((g_fs.RootEntCnt * 32) + g_fs.BytsPerSec - 1) / g_fs.BytsPerSec;
            DirLBA = g_fs.PartLBA + g_fs.RsvdSecCnt + g_fs.NumFATs * g_fs.FATSz;
            DirNSecs = RootSecs;
            DirBufBytes = RootSecs * g_fs.BytsPerSec;
        } else {
            DirLBA = clus_to_lba(&g_fs, CurClus);
            DirNSecs = g_fs.SecPerClus;
            DirBufBytes = g_fs.ClusSize;
        }

        DirBuf = kmalloc(DirBufBytes);
        if (!DirBuf) return EFI_OUT_OF_RESOURCES;

        EFI_STATUS rst = blk_read(g_fs.Dev, DirLBA, DirNSecs, DirBuf);
        if (EFI_ERROR(rst)) { kfree(DirBuf); return rst; }

        dirent_update_size(DirBuf, FileDirOff, NewSize, FileFirst,
                            FileWasEmpty, g_fs.IsFAT32);

        EFI_STATUS wst = blk_write(g_fs.Dev, DirLBA, DirNSecs, DirBuf);
        kfree(DirBuf);
        if (EFI_ERROR(wst)) return wst;
    }

    // --- Flush FAT ---
    return fat_flush(&g_fs);
}