use alloc::{sync::Arc, vec, vec::Vec};
use bytemuck::{Pod, Zeroable};
use yggdrasil_abi::error::Error;

use crate::device::block::BlockDevice;

use super::Partition;

#[derive(Clone, Copy, Debug, Zeroable, Pod)]
#[repr(C)]
struct PartitionEntry {
    attributes: u8,
    start_chs: [u8; 3],
    ty: u8,
    last_chs: [u8; 3],
    start_lba: u32,
    sector_count: u32,
}

pub(crate) async fn probe_mbr(dev: &Arc<dyn BlockDevice>) -> Result<Option<Vec<Partition>>, Error> {
    let mut mbr = [0; 512];
    let mut entries = [PartitionEntry::zeroed(); 4];
    dev.read_exact(0, &mut mbr).await?;

    if mbr[510] != 0x55 || mbr[511] != 0xAA {
        return Ok(None);
    }

    for (i, entry) in entries[0..4].iter_mut().enumerate() {
        let s = 446 + i * size_of::<PartitionEntry>();
        *entry = bytemuck::pod_read_unaligned(&mbr[s..s + size_of::<PartitionEntry>()]);
    }

    let mut partitions = vec![];

    for entry in entries {
        if entry.start_lba != 0
            && entry.sector_count != 0
            && let Some(end) = entry.start_lba.checked_add(entry.sector_count)
        {
            let partition = Partition::new(dev.clone(), entry.start_lba.into(), end.into());
            partitions.push(partition);
        }
    }

    Ok(Some(partitions))
}
