/**
 * @file mkfs.fat32.c
 * @brief Format a block device with a FAT32 filesystem.
 *
 * Writes the FAT32 boot sector (BPB), FSInfo sector, two File Allocation
 * Tables, and an empty root directory cluster to the specified device.
 *
 * Usage: mkfs.fat32 <device> <total_sectors>
 *   e.g.: mkfs.fat32 hdd1mbr1 14336
 *
 * The device is accessed via /dev/<device>. Total sectors must be
 * specified because the OS does not yet expose device size to userspace.
 */

#include "syscalls.h"

typedef unsigned char uint8_t;
typedef unsigned short uint16_t;

/* ================================================================
 * Constants
 * ================================================================ */

#define SECTOR_SIZE     512
#define RESERVED_SECTS  32
#define NUM_FATS        2
#define ROOT_CLUSTER    2
#define FSINFO_SECTOR   1
#define BACKUP_BOOT     6

/* ================================================================
 * Helpers
 * ================================================================ */

/**
 * Print a 32-bit decimal value.
 */
static void print_dec(uint32_t val) {
    char buf[12];
    int pos = 11;
    buf[pos] = '\0';
    if (val == 0) {
        buf[--pos] = '0';
    } else {
        while (val > 0 && pos > 0) {
            buf[--pos] = (char)('0' + (val % 10));
            val /= 10;
        }
    }
    puts(&buf[pos]);
}

/**
 * Print a hex value.
 */
static void print_hex(uint32_t val) {
    char buf[11];
    buf[0] = '0';
    buf[1] = 'x';
    for (int i = 9; i >= 2; i--) {
        int n = (int)(val & 0xF);
        buf[i] = (n < 10) ? (char)('0' + n) : (char)('A' + n - 10);
        val >>= 4;
    }
    buf[10] = '\0';
    puts(buf);
}

/**
 * Parse decimal string.
 */
static uint32_t parse_dec(const char *s, int *ok) {
    uint32_t val = 0;
    *ok = 0;
    if (!s || !*s) return 0;
    while (*s) {
        if (*s < '0' || *s > '9') { *ok = 0; return 0; }
        val = val * 10 + (uint32_t)(*s - '0');
        s++;
    }
    *ok = 1;
    return val;
}

/**
 * Write a 16-bit little-endian value.
 */
static void w16(uint8_t *p, uint32_t val) {
    p[0] = (uint8_t)(val & 0xFF);
    p[1] = (uint8_t)((val >> 8) & 0xFF);
}

/**
 * Write a 32-bit little-endian value.
 */
static void w32(uint8_t *p, uint32_t val) {
    p[0] = (uint8_t)(val & 0xFF);
    p[1] = (uint8_t)((val >> 8) & 0xFF);
    p[2] = (uint8_t)((val >> 16) & 0xFF);
    p[3] = (uint8_t)((val >> 24) & 0xFF);
}

/* ================================================================
 * Argument parsing
 * ================================================================ */

#define MAX_TOKENS 4
static char tok_buf[256];
static char *tokens[MAX_TOKENS];
static int token_count;

static void tokenize(const char *input) {
    int i = 0;
    while (input[i] && i < 255) { tok_buf[i] = input[i]; i++; }
    tok_buf[i] = '\0';
    token_count = 0;
    char *p = tok_buf;
    while (*p && token_count < MAX_TOKENS) {
        while (*p == ' ') p++;
        if (!*p) break;
        tokens[token_count++] = p;
        while (*p && *p != ' ') p++;
        if (*p) { *p = '\0'; p++; }
    }
}

/* ================================================================
 * Sector buffer (reused for each write)
 * ================================================================ */

static uint8_t sector[SECTOR_SIZE];

/**
 * Write the sector buffer to the device fd.
 * Returns 0 on success, -1 on failure.
 */
static int write_sector(int32_t fd) {
    int32_t n = write(fd, sector, SECTOR_SIZE);
    return (n == SECTOR_SIZE) ? 0 : -1;
}

/**
 * Write N zero sectors to the device fd.
 */
static int write_zero_sectors(int32_t fd, uint32_t count) {
    memset(sector, 0, SECTOR_SIZE);
    for (uint32_t i = 0; i < count; i++) {
        if (write_sector(fd) != 0) return -1;
    }
    return 0;
}

/* ================================================================
 * FAT32 geometry calculation
 * ================================================================ */

/**
 * Determine sectors per cluster based on volume size.
 * Uses Microsoft's recommended values for FAT32.
 */
static uint32_t choose_spc(uint32_t total_sectors) {
    /* For small volumes, use small clusters to maximize cluster count.
     * FAT32 "officially" needs >= 65525 clusters, but our driver
     * doesn't enforce this minimum. */
    if (total_sectors <= 16384)       return 1;   /* <= 8 MiB: 512B */
    if (total_sectors <= 131072)      return 8;   /* <= 64 MiB: 4K */
    if (total_sectors <= 524288)      return 8;   /* <= 256 MiB: 4K */
    if (total_sectors <= 16777216)    return 8;   /* <= 8 GiB: 4K */
    if (total_sectors <= 33554432)    return 16;  /* <= 16 GiB: 8K */
    return 32;                                     /* > 16 GiB: 16K */
}

/**
 * Calculate FAT size in sectors using Microsoft's formula.
 *
 * fat_size = ceil((total - reserved) / (entries_per_sec * spc + num_fats))
 * where entries_per_sec = 128 (512 / 4 bytes per FAT32 entry).
 */
static uint32_t calc_fat_size(uint32_t total_sectors, uint32_t spc) {
    uint32_t data_area = total_sectors - RESERVED_SECTS;
    uint32_t tmp = 128 * spc + NUM_FATS;
    return (data_area + tmp - 1) / tmp;
}

/* ================================================================
 * BPB / Boot sector construction
 * ================================================================ */

/**
 * Build the FAT32 boot sector (BPB) into the sector buffer.
 */
static void build_boot_sector(uint32_t total_sectors, uint32_t spc,
                               uint32_t fat_size, const char *label) {
    memset(sector, 0, SECTOR_SIZE);

    /* Jump instruction: EB 58 90 (jump over BPB to boot code area) */
    sector[0] = 0xEB;
    sector[1] = 0x58;
    sector[2] = 0x90;

    /* OEM Name (offset 3, 8 bytes) */
    memcpy(&sector[3], "CLAUDEOS", 8);

    /* BPB — BIOS Parameter Block */
    w16(&sector[11], SECTOR_SIZE);       /* bytes_per_sector */
    sector[13] = (uint8_t)spc;          /* sectors_per_cluster */
    w16(&sector[14], RESERVED_SECTS);    /* reserved_sectors */
    sector[16] = NUM_FATS;               /* num_fats */
    w16(&sector[17], 0);                 /* root_entry_count (0 for FAT32) */
    w16(&sector[19], 0);                 /* total_sectors_16 (0, use 32-bit) */
    sector[21] = 0xF8;                   /* media_type (hard disk) */
    w16(&sector[22], 0);                 /* fat_size_16 (0, use 32-bit) */
    w16(&sector[24], 63);               /* sectors_per_track */
    w16(&sector[26], 255);              /* num_heads */
    w32(&sector[28], 0);                /* hidden_sectors */
    w32(&sector[32], total_sectors);    /* total_sectors_32 */

    /* FAT32 Extended BPB (offset 36) */
    w32(&sector[36], fat_size);         /* fat_size_32 */
    w16(&sector[40], 0);                /* ext_flags */
    w16(&sector[42], 0);                /* fs_version */
    w32(&sector[44], ROOT_CLUSTER);     /* root_cluster */
    w16(&sector[48], FSINFO_SECTOR);    /* fs_info sector */
    w16(&sector[50], BACKUP_BOOT);      /* backup_boot sector */
    /* Reserved (12 bytes at offset 52, already zeroed) */

    sector[64] = 0x80;                  /* drive_number (HDD) */
    sector[65] = 0;                     /* reserved1 */
    sector[66] = 0x29;                  /* boot_sig (extended boot sig) */

    /* Volume serial number (we'll use a simple hash) */
    uint32_t serial = total_sectors ^ (spc << 16) ^ 0xC1A0DE05;
    w32(&sector[67], serial);

    /* Volume label (11 bytes, space-padded) */
    int li = 0;
    while (label[li] && li < 11) {
        char c = label[li];
        /* Convert to uppercase for FAT compatibility */
        if (c >= 'a' && c <= 'z') c = (char)(c - 32);
        sector[71 + li] = (uint8_t)c;
        li++;
    }
    while (li < 11) {
        sector[71 + li] = ' ';
        li++;
    }

    /* FS type string (8 bytes) */
    memcpy(&sector[82], "FAT32   ", 8);

    /* Boot sector signature */
    sector[510] = 0x55;
    sector[511] = 0xAA;
}

/* ================================================================
 * FSInfo sector construction
 * ================================================================ */

/**
 * Build the FSInfo sector into the sector buffer.
 */
static void build_fsinfo(uint32_t free_clusters, uint32_t next_free) {
    memset(sector, 0, SECTOR_SIZE);

    /* FSInfo signatures */
    w32(&sector[0], 0x41615252);        /* Lead signature */
    w32(&sector[484], 0x61417272);      /* Structure signature */
    w32(&sector[488], free_clusters);   /* Free cluster count */
    w32(&sector[492], next_free);       /* Next free cluster hint */
    w32(&sector[508], 0xAA550000);      /* Trail signature */
}

/* ================================================================
 * FAT table construction
 * ================================================================ */

/**
 * Build the first sector of the FAT into the sector buffer.
 * Entry 0 = media byte marker, Entry 1 = EOC, Entry 2 = EOC (root dir).
 */
static void build_fat_first_sector(void) {
    memset(sector, 0, SECTOR_SIZE);

    /* FAT entry 0: media type in low byte, rest 0xFF */
    w32(&sector[0], 0x0FFFFFF8);

    /* FAT entry 1: end-of-chain marker */
    w32(&sector[4], 0x0FFFFFFF);

    /* FAT entry 2: root directory cluster (end-of-chain) */
    w32(&sector[8], 0x0FFFFFFF);
}

/* ================================================================
 * Root directory construction
 * ================================================================ */

/**
 * Build the root directory cluster's first sector.
 * Contains a volume label entry only.
 */
static void build_root_dir(const char *label) {
    memset(sector, 0, SECTOR_SIZE);

    /* Volume label directory entry (32 bytes) */
    int li = 0;
    while (label[li] && li < 11) {
        char c = label[li];
        if (c >= 'a' && c <= 'z') c = (char)(c - 32);
        sector[li] = (uint8_t)c;
        li++;
    }
    while (li < 11) {
        sector[li] = ' ';
        li++;
    }
    sector[11] = 0x08; /* ATTR_VOLUME_ID */
}

/* ================================================================
 * Main formatting routine
 * ================================================================ */

/**
 * Format the device with FAT32.
 *
 * Writes sectors sequentially from sector 0:
 *   [0]       Boot sector (BPB)
 *   [1]       FSInfo
 *   [2..5]    Zero (reserved)
 *   [6]       Backup boot sector
 *   [7]       Backup FSInfo
 *   [8..31]   Zero (reserved)
 *   [32..32+fat_size-1]  FAT 1
 *   [32+fat_size..32+2*fat_size-1]  FAT 2
 *   [data_start..data_start+spc-1]  Root directory cluster
 */
static int format_device(const char *dev_path, uint32_t total_sectors,
                          const char *label) {
    uint32_t spc = choose_spc(total_sectors);
    uint32_t fat_size = calc_fat_size(total_sectors, spc);
    uint32_t data_start = RESERVED_SECTS + NUM_FATS * fat_size;
    uint32_t data_sectors = total_sectors - data_start;
    uint32_t total_clusters = data_sectors / spc;
    uint32_t free_clusters = total_clusters - 1; /* Root dir uses cluster 2 */

    /* Print format parameters */
    puts("Formatting ");
    puts(dev_path);
    puts(" as FAT32\n");
    puts("  Total sectors:      ");
    print_dec(total_sectors);
    puts("\n  Sectors/cluster:    ");
    print_dec(spc);
    puts("\n  Reserved sectors:   ");
    print_dec(RESERVED_SECTS);
    puts("\n  FAT size (sectors): ");
    print_dec(fat_size);
    puts("\n  Number of FATs:     ");
    print_dec(NUM_FATS);
    puts("\n  Data start sector:  ");
    print_dec(data_start);
    puts("\n  Total clusters:     ");
    print_dec(total_clusters);
    puts("\n  Volume label:       ");
    puts(label);
    puts("\n");

    /* Sanity checks */
    if (total_sectors < 128) {
        puts("Error: volume too small for FAT32\n");
        return -1;
    }
    if (data_start >= total_sectors) {
        puts("Error: FAT tables don't fit in volume\n");
        return -1;
    }

    /* Open device for writing */
    int32_t fd = open(dev_path, 0);
    if (fd < 0) {
        puts("Error: cannot open ");
        puts(dev_path);
        puts("\n");
        return -1;
    }

    puts("Writing boot sector...\n");

    /* Sector 0: Boot sector */
    build_boot_sector(total_sectors, spc, fat_size, label);
    if (write_sector(fd) != 0) goto fail;

    /* Sector 1: FSInfo */
    build_fsinfo(free_clusters, 3); /* Next free = cluster 3 */
    if (write_sector(fd) != 0) goto fail;

    /* Sectors 2-5: zero */
    if (write_zero_sectors(fd, 4) != 0) goto fail;

    /* Sector 6: Backup boot sector */
    build_boot_sector(total_sectors, spc, fat_size, label);
    if (write_sector(fd) != 0) goto fail;

    /* Sector 7: Backup FSInfo */
    build_fsinfo(free_clusters, 3);
    if (write_sector(fd) != 0) goto fail;

    /* Sectors 8-31: zero (remaining reserved area) */
    if (write_zero_sectors(fd, RESERVED_SECTS - 8) != 0) goto fail;

    puts("Writing FAT 1...\n");

    /* FAT 1: first sector has entries 0-2 */
    build_fat_first_sector();
    if (write_sector(fd) != 0) goto fail;

    /* FAT 1: remaining sectors (all zero = free) */
    if (fat_size > 1) {
        if (write_zero_sectors(fd, fat_size - 1) != 0) goto fail;
    }

    puts("Writing FAT 2...\n");

    /* FAT 2: identical copy */
    build_fat_first_sector();
    if (write_sector(fd) != 0) goto fail;
    if (fat_size > 1) {
        if (write_zero_sectors(fd, fat_size - 1) != 0) goto fail;
    }

    puts("Writing root directory...\n");

    /* Root directory cluster (first sector has volume label) */
    build_root_dir(label);
    if (write_sector(fd) != 0) goto fail;

    /* Remaining sectors in root cluster (zero) */
    if (spc > 1) {
        if (write_zero_sectors(fd, spc - 1) != 0) goto fail;
    }

    close(fd);

    puts("Format complete!\n");
    puts("  ");
    print_dec(data_start + spc);
    puts(" sectors written\n");

    return 0;

fail:
    close(fd);
    puts("Error: write failed!\n");
    return -1;
}

/* ================================================================
 * Main
 * ================================================================ */

int main(void) {
    char arg1[256];
    if (getenv("ARG1", arg1, sizeof(arg1)) < 0 || arg1[0] == '\0') {
        puts("mkfs.fat32 - Format a device with FAT32\n");
        puts("\n");
        puts("Usage: mkfs.fat32 <device> <total_sectors> [label]\n");
        puts("\n");
        puts("  device:        Device name (e.g., hdd1mbr1)\n");
        puts("  total_sectors: Number of 512-byte sectors\n");
        puts("  label:         Volume label (default: CLAUDEOS)\n");
        puts("\n");
        puts("Example: mkfs.fat32 hdd1mbr1 14336\n");
        return 1;
    }

    tokenize(arg1);
    if (token_count < 2) {
        puts("mkfs.fat32: need <device> and <total_sectors>\n");
        puts("  Use 'diskpart <dev> list' to find sector count.\n");
        return 1;
    }

    /* Build /dev/<device> path */
    char dev_path[64];
    strcpy(dev_path, "/dev/");
    char *dp = dev_path + 5;
    char *dn = tokens[0];
    while (*dn && dp < dev_path + 62) *dp++ = *dn++;
    *dp = '\0';

    /* Parse total sectors */
    int ok;
    uint32_t total_sectors = parse_dec(tokens[1], &ok);
    if (!ok || total_sectors == 0) {
        puts("mkfs.fat32: invalid sector count '");
        puts(tokens[1]);
        puts("'\n");
        return 1;
    }

    /* Volume label (default "CLAUDEOS") */
    const char *label = "CLAUDEOS";
    if (token_count >= 3) {
        label = tokens[2];
    }

    return (format_device(dev_path, total_sectors, label) == 0) ? 0 : 1;
}