//
// Created by rick on 06-02-21.
//

#include "mbr.h"
#include <types.h>
#include <drivers/ide.h>
#include <mem/mem.h>
#include <libc/kprintf.h>
#include <fs/blockdev.h>

typedef struct {
    u8 bootable;
    u8 starting_head;
    u8 starting_sector: 6;
    u16 starting_cylinder: 10;
    u8 system_id;
    u8 ending_head;
    u8 ending_sector: 6;
    u16 ending_cylinder: 10;
    u32 start_lba;
    u32 num_lbas;
}  __attribute((packed)) mbr_partition_table_entry;

typedef struct {
    u32 unique_id;
    u16 reserved;
    mbr_partition_table_entry entries[4];
    u8 signature[2];
} __attribute((packed)) mbr_table;

typedef struct {
    const block_device *device;
    u32 start_lba;
} mbr_block_driver_info;

u8 mbr_block_dev_access(const block_device *device, u8 direction, u32 lba, u8 sectors, void *target) {
    if (!device->flags.present || lba > device->num_lba) {
        return BLOCK_DEV_ACCESS_ERR;
    }
    const mbr_block_driver_info *info = device->device_info;
    const u32 actual_lba = info->start_lba + lba;
    if (actual_lba > info->device->num_lba) {
        return BLOCK_DEV_ACCESS_ERR;
    }

    // delegate to backing driver
    return info->device->access(info->device, direction, actual_lba, sectors, target);
}

u8 mbr_check_device(const block_device *device, u8 *first_sector) {
    mbr_table *table = (mbr_table *) &first_sector[512 - sizeof(mbr_table)];
    if (table->signature[0] != 0x55 && table->signature[1] != 0xAA) { // AA 55 but in little endian
        return BLOCK_DEV_DRIVER_CHECK_NO_MATCH;
    }

    if (table->reserved != 0x0000 && table->reserved != 0x5A5A) {
        return BLOCK_DEV_DRIVER_CHECK_NO_MATCH;
    }

    for (int i = 0; i < 4; ++i) {
        if (table->entries[i].num_lbas > device->num_lba || table->entries[i].start_lba > device->num_lba) {
            return BLOCK_DEV_DRIVER_CHECK_NO_MATCH;
        }
    }

    for (int i = 0; i < 4; ++i) {
        if (table->entries[i].system_id == 0) continue;

        mbr_block_driver_info *info = malloc(sizeof(mbr_block_driver_info));
        info->device = device;
        info->start_lba = table->entries[i].start_lba;

        block_device logical_device = {
                .flags.present = 1,

                .num_lba = table->entries[i].num_lbas,
                .block_size = device->block_size,
                .access = mbr_block_dev_access,
                .device_info = info,
        };

        sprintf(logical_device.identifier, "%sp%d", device->identifier, i);

        block_dev_register(&logical_device);

    }

    return BLOCK_DEV_DRIVER_CHECK_OK;
}

block_dev_driver mbr_driver = {
        .name = "mbr",
        .flags.root_only = 1,
        .check_device = mbr_check_device,
        .free_device = NULL, // todo
};

void mbr_register_block_driver() {
    block_dev_register_driver(&mbr_driver);
}

void mbr_read_from_ide(u8 ide_drive) {
    char *mbr_data = malloc(0x200);
    ide_access(0, ide_drive, 0, 1, mbr_data);
    mbr_partition_table_entry *entry = (mbr_partition_table_entry *) (mbr_data + 0x1BE);
    printf("Start at %d, count: %d\n", entry->start_lba, entry->num_lbas);
    free(mbr_data);
}