/**
 * @file mbr.c
 * @brief MBR partition table driver implementation.
 *
 * Reads sector 0 from hard drives registered with the devicefs, parses
 * the MBR partition table, and registers each partition as a sub-device.
 * Partition devices are named using the parent device name + "mbr" class
 * prefix, with the number assigned by the devicefs (e.g., hdd1mbr1).
 *
 * Each partition sub-device translates LBA offsets relative to its own
 * start into absolute LBA addresses on the parent device.
 */

#include "mbr.h"
#include "devicefs.h"
#include <string.h>

/* Debug print helpers defined in kernel.c */
extern void offset_print(const char *str);
extern void print_hex(uint32_t val);

/** Maximum number of MBR partition sub-devices. */
#define MBR_MAX_SUBS 16

/**
 * Per-partition context for sub-device read/write operations.
 */
typedef struct mbr_sub_device {
    devicefs_device_t *parent;       /**< Parent block device. */
    uint32_t lba_start;              /**< Partition start LBA on parent. */
    uint32_t sector_count;           /**< Partition size in sectors. */
    uint8_t type;                    /**< MBR partition type code. */
    int active;                      /**< 1 if in use. */
} mbr_sub_device_t;

/** Pool of sub-device contexts. */
static mbr_sub_device_t sub_devices[MBR_MAX_SUBS];

/**
 * Allocate a sub-device context.
 * @return Pointer to a free sub-device, or NULL if exhausted.
 */
static mbr_sub_device_t *alloc_sub(void) {
    for (int i = 0; i < MBR_MAX_SUBS; i++) {
        if (!sub_devices[i].active) {
            sub_devices[i].active = 1;
            return &sub_devices[i];
        }
    }
    return NULL;
}

/* ================================================================
 * Block device ops for MBR partition sub-devices
 * ================================================================ */

/**
 * Read sectors from a partition (translates LBA to parent device).
 */
static int mbr_sub_read(void *dev_data, uint32_t lba,
                        uint32_t count, void *buf) {
    mbr_sub_device_t *sub = (mbr_sub_device_t *)dev_data;
    if (!sub || !sub->parent) return -1;

    /* Bounds check */
    if (lba + count > sub->sector_count) return -1;

    /* Translate to absolute LBA on parent device */
    uint32_t abs_lba = sub->lba_start + lba;

    /* Read from parent device */
    if (!sub->parent->block_ops || !sub->parent->block_ops->read_sectors) {
        return -1;
    }
    return sub->parent->block_ops->read_sectors(sub->parent->dev_data,
                                                 abs_lba, count, buf);
}

/**
 * Write sectors to a partition (translates LBA to parent device).
 */
static int mbr_sub_write(void *dev_data, uint32_t lba,
                         uint32_t count, const void *buf) {
    mbr_sub_device_t *sub = (mbr_sub_device_t *)dev_data;
    if (!sub || !sub->parent) return -1;

    if (lba + count > sub->sector_count) return -1;

    uint32_t abs_lba = sub->lba_start + lba;

    if (!sub->parent->block_ops || !sub->parent->block_ops->write_sectors) {
        return -1;
    }
    return sub->parent->block_ops->write_sectors(sub->parent->dev_data,
                                                  abs_lba, count, buf);
}

/**
 * Get sector size (same as parent device).
 */
static uint32_t mbr_sub_sector_size(void *dev_data) {
    mbr_sub_device_t *sub = (mbr_sub_device_t *)dev_data;
    if (!sub || !sub->parent || !sub->parent->block_ops) return 512;
    if (sub->parent->block_ops->sector_size) {
        return sub->parent->block_ops->sector_size(sub->parent->dev_data);
    }
    return 512;
}

/**
 * Get total sector count of the partition.
 */
static uint32_t mbr_sub_sector_count(void *dev_data) {
    mbr_sub_device_t *sub = (mbr_sub_device_t *)dev_data;
    return sub ? sub->sector_count : 0;
}

/** Block operations for MBR partition sub-devices. */
static devicefs_block_ops_t mbr_sub_ops = {
    .read_sectors  = mbr_sub_read,
    .write_sectors = mbr_sub_write,
    .sector_size   = mbr_sub_sector_size,
    .sector_count  = mbr_sub_sector_count,
};

/* ================================================================
 * MBR scanning
 * ================================================================ */

int mbr_scan(const char *parent_name) {
    devicefs_device_t *parent = devicefs_find(parent_name);
    if (!parent) {
        offset_print("  MBR: device not found: ");
        offset_print(parent_name);
        offset_print("\n");
        return -1;
    }

    if (parent->type != DEVICEFS_BLOCK || !parent->block_ops) {
        offset_print("  MBR: not a block device: ");
        offset_print(parent_name);
        offset_print("\n");
        return -1;
    }

    if (!parent->block_ops->read_sectors) {
        offset_print("  MBR: device has no read_sectors op\n");
        return -1;
    }

    /* Read sector 0 (MBR) */
    uint8_t sector[512];
    if (parent->block_ops->read_sectors(parent->dev_data, 0, 1, sector) != 0) {
        offset_print("  MBR: failed to read sector 0 of ");
        offset_print(parent_name);
        offset_print("\n");
        return -1;
    }

    /* Check MBR signature */
    uint16_t sig = (uint16_t)sector[510] | ((uint16_t)sector[511] << 8);
    if (sig != MBR_SIGNATURE) {
        offset_print("  MBR: no valid signature on ");
        offset_print(parent_name);
        offset_print(" (sig=");
        print_hex(sig);
        offset_print(")\n");
        return 0; /* Not an error, just no MBR */
    }

    /* Build the class name for partition sub-devices: parentName + "mbr" */
    char class_name[DEVICEFS_MAX_DEV_NAME];
    memset(class_name, 0, sizeof(class_name));
    strncpy(class_name, parent_name, DEVICEFS_MAX_DEV_NAME - 4);
    /* Append "mbr" */
    uint32_t clen = strlen(class_name);
    if (clen + 3 < DEVICEFS_MAX_DEV_NAME) {
        class_name[clen] = 'm';
        class_name[clen + 1] = 'b';
        class_name[clen + 2] = 'r';
        class_name[clen + 3] = '\0';
    }

    /* Parse the 4 partition table entries (at offset 446) */
    int found = 0;
    for (int i = 0; i < MBR_MAX_PARTITIONS; i++) {
        mbr_partition_entry_t *entry =
            (mbr_partition_entry_t *)(sector + 446 + i * 16);

        /* Skip empty partitions */
        if (entry->type == MBR_TYPE_EMPTY) continue;
        if (entry->sector_count == 0) continue;

        /* Skip extended partitions for now */
        if (entry->type == MBR_TYPE_EXTENDED) {
            offset_print("  MBR: skipping extended partition on ");
            offset_print(parent_name);
            offset_print("\n");
            continue;
        }

        mbr_sub_device_t *sub = alloc_sub();
        if (!sub) {
            offset_print("  MBR: no free sub-device slots\n");
            break;
        }

        sub->parent = parent;
        sub->lba_start = entry->lba_start;
        sub->sector_count = entry->sector_count;
        sub->type = entry->type;

        /* Register with devicefs */
        devicefs_device_t *dev = devicefs_register_block(class_name,
                                                          &mbr_sub_ops,
                                                          sub);
        if (dev) {
            offset_print("  MBR: ");
            offset_print(dev->name);
            offset_print(" type=");
            print_hex(entry->type);
            offset_print(" LBA=");
            print_hex(entry->lba_start);
            offset_print(" size=");
            print_hex(entry->sector_count);
            offset_print(" sectors\n");
            found++;
        }
    }

    if (found == 0) {
        offset_print("  MBR: no partitions on ");
        offset_print(parent_name);
        offset_print("\n");
    }

    return found;
}

int init_mbr(void) {
    memset(sub_devices, 0, sizeof(sub_devices));

    /* Scan all "hdd" class devices for MBR partitions */
    int total_partitions = 0;
    for (uint32_t i = 1; i <= 16; i++) {
        /* Build device name: hdd1, hdd2, ... */
        char name[DEVICEFS_MAX_DEV_NAME];
        memset(name, 0, sizeof(name));
        name[0] = 'h'; name[1] = 'd'; name[2] = 'd';
        /* Append number */
        if (i < 10) {
            name[3] = (char)('0' + i);
        } else {
            name[3] = (char)('0' + i / 10);
            name[4] = (char)('0' + i % 10);
        }

        devicefs_device_t *dev = devicefs_find(name);
        if (!dev) continue;

        int n = mbr_scan(name);
        if (n > 0) total_partitions += n;
    }

    offset_print("  MBR: ");
    print_hex((uint32_t)total_partitions);
    offset_print(" partition(s) found\n");

    return 0;
}