/**
 * @file diskpart.c
 * @brief MBR partition table editor.
 *
 * Allows viewing, creating, and deleting MBR partitions on block devices.
 *
 * Usage: diskpart <device> [command] [args...]
 *
 * Commands:
 *   list                                    Show partition table (default)
 *   create <slot> <type> <start> <sectors>  Create a partition
 *   delete <slot>                           Delete a partition
 *   activate <slot>                         Set bootable flag
 *   deactivate <slot>                       Clear bootable flag
 *
 * <slot> is 1-4 (MBR supports up to 4 primary partitions).
 * <type> is a hex partition type code (e.g., 0C for FAT32 LBA).
 * <start> and <sectors> are decimal LBA values.
 *
 * Examples:
 *   diskpart hdd1 list
 *   diskpart hdd1 create 1 0C 2048 14336
 *   diskpart hdd1 delete 2
 *   diskpart hdd1 activate 1
 */

#include "syscalls.h"

typedef unsigned char uint8_t;

/** MBR signature. */
#define MBR_SIG_LO 0x55
#define MBR_SIG_HI 0xAA

/** MBR partition entry offsets. */
#define MBR_PART_TABLE 446
#define MBR_PART_SIZE  16

/* ================================================================
 * Numeric formatting helpers
 * ================================================================ */

/**
 * Print a 32-bit value as a hexadecimal string with "0x" prefix.
 */
static void print_hex(uint32_t val) {
    char buf[11]; /* "0x" + 8 digits + NUL */
    buf[0] = '0';
    buf[1] = 'x';
    for (int i = 9; i >= 2; i--) {
        int nibble = (int)(val & 0xF);
        buf[i] = (nibble < 10) ? (char)('0' + nibble)
                               : (char)('A' + nibble - 10);
        val >>= 4;
    }
    buf[10] = '\0';
    puts(buf);
}

/**
 * Print a 32-bit value as a decimal string.
 */
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]);
}

/**
 * Parse a decimal string to uint32_t.
 * Returns the parsed value; sets *ok to 1 on success, 0 on failure.
 */
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;
}

/**
 * Parse a hexadecimal string (without "0x" prefix) to uint32_t.
 */
static uint32_t parse_hex(const char *s, int *ok) {
    uint32_t val = 0;
    *ok = 0;
    if (!s || !*s) return 0;

    /* Skip optional "0x" prefix */
    if (s[0] == '0' && (s[1] == 'x' || s[1] == 'X')) {
        s += 2;
    }
    if (!*s) return 0;

    while (*s) {
        uint32_t digit;
        if (*s >= '0' && *s <= '9')      digit = (uint32_t)(*s - '0');
        else if (*s >= 'a' && *s <= 'f') digit = (uint32_t)(*s - 'a' + 10);
        else if (*s >= 'A' && *s <= 'F') digit = (uint32_t)(*s - 'A' + 10);
        else { *ok = 0; return 0; }
        val = (val << 4) | digit;
        s++;
    }
    *ok = 1;
    return val;
}

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

/** Maximum number of tokens we parse from ARG1. */
#define MAX_TOKENS 8

/** Token buffer storage. */
static char token_buf[256];
static char *tokens[MAX_TOKENS];
static int token_count;

/**
 * Split the ARG1 string into space-separated tokens.
 */
static void tokenize(const char *input) {
    /* Copy input to mutable buffer */
    int i = 0;
    while (input[i] && i < 255) {
        token_buf[i] = input[i];
        i++;
    }
    token_buf[i] = '\0';

    token_count = 0;
    char *p = token_buf;
    while (*p && token_count < MAX_TOKENS) {
        /* Skip whitespace */
        while (*p == ' ') p++;
        if (!*p) break;

        tokens[token_count++] = p;

        /* Advance to next space or end */
        while (*p && *p != ' ') p++;
        if (*p) {
            *p = '\0';
            p++;
        }
    }
}

/* ================================================================
 * Partition type name lookup
 * ================================================================ */

/**
 * Return a human-readable name for an MBR partition type code.
 */
static const char *type_name(uint32_t type) {
    switch (type) {
    case 0x00: return "Empty";
    case 0x01: return "FAT12";
    case 0x04: return "FAT16 <32M";
    case 0x05: return "Extended";
    case 0x06: return "FAT16 >=32M";
    case 0x07: return "NTFS/HPFS";
    case 0x0B: return "FAT32";
    case 0x0C: return "FAT32 LBA";
    case 0x0E: return "FAT16 LBA";
    case 0x82: return "Linux swap";
    case 0x83: return "Linux";
    default:   return "Unknown";
    }
}

/* ================================================================
 * MBR reading / writing
 * ================================================================ */

/** MBR sector buffer (512 bytes). */
static uint8_t mbr[512];

/**
 * Read a 32-bit little-endian value from a byte array.
 */
static uint32_t read32le(const uint8_t *p) {
    return (uint32_t)p[0]
         | ((uint32_t)p[1] << 8)
         | ((uint32_t)p[2] << 16)
         | ((uint32_t)p[3] << 24);
}

/**
 * Write a 32-bit little-endian value to a byte array.
 */
static void write32le(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);
}

/**
 * Read MBR (sector 0) from a device.
 * @param dev_path Full path like "/dev/hdd1".
 * @return 0 on success, -1 on failure.
 */
static int read_mbr(const char *dev_path) {
    int32_t fd = open(dev_path, 0);
    if (fd < 0) {
        puts("diskpart: cannot open ");
        puts(dev_path);
        puts("\n");
        return -1;
    }
    int32_t n = read(fd, mbr, 512);
    close(fd);
    if (n != 512) {
        puts("diskpart: failed to read MBR (got ");
        print_dec((uint32_t)(n < 0 ? 0 : n));
        puts(" bytes)\n");
        return -1;
    }
    return 0;
}

/**
 * Write MBR (sector 0) to a device.
 * @param dev_path Full path like "/dev/hdd1".
 * @return 0 on success, -1 on failure.
 */
static int write_mbr(const char *dev_path) {
    int32_t fd = open(dev_path, 0);
    if (fd < 0) {
        puts("diskpart: cannot open ");
        puts(dev_path);
        puts(" for writing\n");
        return -1;
    }
    int32_t n = write(fd, mbr, 512);
    close(fd);
    if (n != 512) {
        puts("diskpart: failed to write MBR\n");
        return -1;
    }
    return 0;
}

/* ================================================================
 * Commands
 * ================================================================ */

/**
 * Display the partition table.
 */
static void cmd_list(void) {
    /* Check signature */
    if (mbr[510] != MBR_SIG_LO || mbr[511] != MBR_SIG_HI) {
        puts("No valid MBR signature found.\n");
        puts("Use 'create' to initialize partitions (signature will be added).\n");
        return;
    }

    puts("Slot  Boot  Type        Start LBA   Sectors     Size\n");
    puts("----  ----  ----------  ----------  ----------  ----------\n");

    for (int i = 0; i < 4; i++) {
        uint8_t *entry = &mbr[MBR_PART_TABLE + i * MBR_PART_SIZE];
        uint8_t status = entry[0];
        uint8_t type = entry[4];
        uint32_t lba_start = read32le(&entry[8]);
        uint32_t sectors = read32le(&entry[12]);

        /* Slot number */
        putchar(' ');
        putchar((char)('1' + i));
        puts("    ");

        /* Boot flag */
        if (status == 0x80)
            puts(" *  ");
        else
            puts("    ");

        /* Type */
        puts("  ");
        print_hex(type);
        puts(" ");

        /* Type name (pad to ~10 chars) */
        const char *tn = type_name(type);
        puts(tn);
        int tnlen = (int)strlen(tn);
        for (int j = tnlen; j < 12; j++) putchar(' ');

        /* Start LBA */
        print_dec(lba_start);
        /* Pad */
        char tmp[12];
        int tpos = 11;
        tmp[tpos] = '\0';
        uint32_t v = lba_start;
        if (v == 0) { tmp[--tpos] = '0'; }
        else { while (v > 0) { tmp[--tpos] = (char)('0' + v % 10); v /= 10; } }
        for (int j = 11 - tpos; j < 12; j++) putchar(' ');

        /* Sectors */
        print_dec(sectors);
        tpos = 11;
        tmp[tpos] = '\0';
        v = sectors;
        if (v == 0) { tmp[--tpos] = '0'; }
        else { while (v > 0) { tmp[--tpos] = (char)('0' + v % 10); v /= 10; } }
        for (int j = 11 - tpos; j < 12; j++) putchar(' ');

        /* Size in KiB or MiB */
        uint32_t kib = sectors / 2;
        if (kib >= 1024) {
            print_dec(kib / 1024);
            puts(" MiB");
        } else {
            print_dec(kib);
            puts(" KiB");
        }

        puts("\n");
    }
}

/**
 * Create a partition.
 * Args: <slot 1-4> <type hex> <start_lba> <sectors>
 */
static int cmd_create(void) {
    if (token_count < 6) {
        puts("Usage: diskpart <dev> create <slot> <type> <start> <sectors>\n");
        puts("  slot: 1-4\n");
        puts("  type: hex partition type (e.g., 0C for FAT32 LBA)\n");
        puts("  start: start LBA (decimal)\n");
        puts("  sectors: partition size in sectors (decimal)\n");
        return -1;
    }

    int ok;
    uint32_t slot = parse_dec(tokens[2], &ok);
    if (!ok || slot < 1 || slot > 4) {
        puts("diskpart: invalid slot (must be 1-4)\n");
        return -1;
    }

    uint32_t type = parse_hex(tokens[3], &ok);
    if (!ok || type > 0xFF || type == 0) {
        puts("diskpart: invalid type (must be 01-FF hex)\n");
        return -1;
    }

    uint32_t start = parse_dec(tokens[4], &ok);
    if (!ok) {
        puts("diskpart: invalid start LBA\n");
        return -1;
    }

    uint32_t sectors = parse_dec(tokens[5], &ok);
    if (!ok || sectors == 0) {
        puts("diskpart: invalid sector count\n");
        return -1;
    }

    uint32_t idx = slot - 1;
    uint8_t *entry = &mbr[MBR_PART_TABLE + idx * MBR_PART_SIZE];

    /* Check if slot is already in use */
    if (entry[4] != 0x00) {
        puts("diskpart: slot ");
        print_dec(slot);
        puts(" is already in use (type ");
        print_hex(entry[4]);
        puts("). Delete it first.\n");
        return -1;
    }

    /* Clear the entry */
    memset(entry, 0, MBR_PART_SIZE);

    /* Set type */
    entry[4] = (uint8_t)type;

    /* Set LBA start and sector count */
    write32le(&entry[8], start);
    write32le(&entry[12], sectors);

    /* Ensure MBR signature is present */
    mbr[510] = MBR_SIG_LO;
    mbr[511] = MBR_SIG_HI;

    puts("Created partition ");
    print_dec(slot);
    puts(": type=");
    print_hex(type);
    puts(" (");
    puts(type_name(type));
    puts(") start=");
    print_dec(start);
    puts(" sectors=");
    print_dec(sectors);
    puts("\n");

    return 0;
}

/**
 * Delete a partition.
 */
static int cmd_delete(void) {
    if (token_count < 3) {
        puts("Usage: diskpart <dev> delete <slot>\n");
        return -1;
    }

    int ok;
    uint32_t slot = parse_dec(tokens[2], &ok);
    if (!ok || slot < 1 || slot > 4) {
        puts("diskpart: invalid slot (must be 1-4)\n");
        return -1;
    }

    uint32_t idx = slot - 1;
    uint8_t *entry = &mbr[MBR_PART_TABLE + idx * MBR_PART_SIZE];

    if (entry[4] == 0x00) {
        puts("diskpart: slot ");
        print_dec(slot);
        puts(" is already empty\n");
        return 0;
    }

    uint8_t old_type = entry[4];
    memset(entry, 0, MBR_PART_SIZE);

    puts("Deleted partition ");
    print_dec(slot);
    puts(" (was type ");
    print_hex(old_type);
    puts(" - ");
    puts(type_name(old_type));
    puts(")\n");

    return 0;
}

/**
 * Set or clear the bootable flag on a partition.
 */
static int cmd_boot(int activate) {
    if (token_count < 3) {
        puts("Usage: diskpart <dev> ");
        puts(activate ? "activate" : "deactivate");
        puts(" <slot>\n");
        return -1;
    }

    int ok;
    uint32_t slot = parse_dec(tokens[2], &ok);
    if (!ok || slot < 1 || slot > 4) {
        puts("diskpart: invalid slot (must be 1-4)\n");
        return -1;
    }

    uint32_t idx = slot - 1;
    uint8_t *entry = &mbr[MBR_PART_TABLE + idx * MBR_PART_SIZE];

    if (entry[4] == 0x00) {
        puts("diskpart: slot ");
        print_dec(slot);
        puts(" is empty\n");
        return -1;
    }

    if (activate) {
        /* Clear bootable flag on all other partitions */
        for (int i = 0; i < 4; i++) {
            mbr[MBR_PART_TABLE + i * MBR_PART_SIZE] = 0x00;
        }
        entry[0] = 0x80;
        puts("Partition ");
        print_dec(slot);
        puts(" is now bootable\n");
    } else {
        entry[0] = 0x00;
        puts("Partition ");
        print_dec(slot);
        puts(" is no longer bootable\n");
    }

    return 0;
}

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

int main(void) {
    char arg1[256];
    if (getenv("ARG1", arg1, sizeof(arg1)) < 0 || arg1[0] == '\0') {
        puts("diskpart - MBR partition table editor\n");
        puts("\n");
        puts("Usage: diskpart <device> [command] [args...]\n");
        puts("\n");
        puts("Commands:\n");
        puts("  list                                   Show partitions (default)\n");
        puts("  create <slot> <type> <start> <sectors>  Create partition\n");
        puts("  delete <slot>                           Delete partition\n");
        puts("  activate <slot>                         Set bootable flag\n");
        puts("  deactivate <slot>                       Clear bootable flag\n");
        puts("\n");
        puts("Slot: 1-4, type: hex (e.g., 0C=FAT32 LBA, 83=Linux)\n");
        return 1;
    }

    tokenize(arg1);
    if (token_count < 1) {
        puts("diskpart: no device specified\n");
        return 1;
    }

    /* Token 0 is the device name. Build /dev/<device> path. */
    char dev_path[64];
    strcpy(dev_path, "/dev/");
    /* Append device name */
    char *dp = dev_path + 5;
    char *dn = tokens[0];
    while (*dn && dp < dev_path + 62) {
        *dp++ = *dn++;
    }
    *dp = '\0';

    /* Determine command (default is "list") */
    const char *cmd = "list";
    if (token_count >= 2) {
        cmd = tokens[1];
    }

    /* For "list" we only read; for others we read-modify-write */
    int need_write = 0;

    if (strcmp(cmd, "list") == 0) {
        /* Just read and display */
        if (read_mbr(dev_path) != 0) return 1;
        cmd_list();
        return 0;
    }

    /* Read current MBR for modification commands */
    if (read_mbr(dev_path) != 0) return 1;

    int result = 0;
    if (strcmp(cmd, "create") == 0) {
        result = cmd_create();
        if (result == 0) need_write = 1;
    } else if (strcmp(cmd, "delete") == 0) {
        result = cmd_delete();
        if (result == 0) need_write = 1;
    } else if (strcmp(cmd, "activate") == 0) {
        result = cmd_boot(1);
        if (result == 0) need_write = 1;
    } else if (strcmp(cmd, "deactivate") == 0) {
        result = cmd_boot(0);
        if (result == 0) need_write = 1;
    } else {
        puts("diskpart: unknown command '");
        puts(cmd);
        puts("'\n");
        return 1;
    }

    if (result != 0) return 1;

    /* Write back modified MBR */
    if (need_write) {
        if (write_mbr(dev_path) != 0) {
            puts("diskpart: WARNING - failed to write MBR!\n");
            return 1;
        }
        puts("MBR written successfully.\n");
    }

    return 0;
}