#include "fs/MBR.h"
#include "Log.h"
#include "fs/GPT.h"
#include <luna/CType.h>

static Result<String> create_partition_name(SharedPtr<Device> host_device, u32 partition_index)
{
    auto host_path = host_device->device_path();

    char last = host_path[host_path.length() - 1];

    if (_isdigit(last)) return String::format("%sp%d"_sv, host_path.chars(), partition_index);

    return String::format("%s%d"_sv, host_path.chars(), partition_index);
}

namespace MBR
{
    Result<void> PartitionDevice::create(SharedPtr<Device> host_device, usize start_block, usize num_blocks,
                                         u32 partition_index)
    {
        static u32 next_minor = 0;
        auto device = TRY(adopt_shared_if_nonnull(new (std::nothrow) PartitionDevice()));
        device->m_host_device = host_device;
        device->m_start_offset = start_block * device->m_block_size;
        device->m_num_blocks = num_blocks;
        device->m_device_path = TRY(create_partition_name(host_device, partition_index));
        return DeviceRegistry::register_special_device(DeviceRegistry::DiskPartition, next_minor++, device, 0400);
    }

    Result<usize> PartitionDevice::read(u8* buf, usize offset, usize length) const
    {
        if (length == 0) return 0;

        if (offset > size()) return 0;
        if (offset + length > size()) length = size() - offset;

        return m_host_device->read(buf, m_start_offset + offset, length);
    }

    Result<usize> PartitionDevice::write(const u8* buf, usize offset, usize length)
    {
        if (length == 0) return 0;

        if (offset > size()) return 0;
        if (offset + length > size()) length = size() - offset;

        return m_host_device->write(buf, m_start_offset + offset, length);
    }

    Result<bool> identify(SharedPtr<Device> device)
    {
        // Cannot read a partition table from a character device! Who is even coming up with this silliness?
        if (!device->is_block_device()) return false;

        DiskHeader hdr;
        const usize nread = TRY(device->read((u8*)&hdr, 0, sizeof(hdr)));
        check(nread == 512);

        if (hdr.signature[0] != MBR_SIGNATURE_1 || hdr.signature[1] != MBR_SIGNATURE_2) return false;

        u32 partition_index = 1;

        for (int i = 0; i < 4; i++)
        {
            const auto& part = hdr.partitions[i];
            if (part.partition_type == 0xee) return GPT::identify(device);
        }

        for (int i = 0; i < 4; i++)
        {
            const auto& part = hdr.partitions[i];
            if (part.partition_type == 0) continue; // Not active.

            bool bootable = part.attributes & MBR_BOOTABLE;
            kinfoln("mbr: Partition #%d is active: bootable=%d, type=%x, start=%d, sectors=%d", i, bootable,
                    part.partition_type, part.start_lba, part.num_sectors);

            TRY(PartitionDevice::create(device, part.start_lba, part.num_sectors, partition_index++));
        }

        return true;
    }
}