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

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

use super::Partition;

#[derive(Clone, Copy, Pod, Zeroable)]
#[repr(C)]
struct GptHeader {
    signature: [u8; 8],
    revision: u32,
    header_size: u32,
    crc32: u32,
    _0: u32,
    header_lba: u64,
    alternate_header_lba: u64,
    first_usable_lba: u64,
    last_usable_lba: u64,
    guid: [u8; 16],
    partition_table_lba: u64,
    partition_table_len: u32,
    partition_table_entry_size: u32,
    partition_table_crc32: u32,
    _1: u32,
}

#[derive(Clone, Copy, Zeroable, Pod)]
#[repr(C)]
struct GptEntry {
    type_guid: Uuid,
    part_guid: Uuid,
    lba_start: u64,
    lba_end: u64,
    attrs: u64,
}

async unsafe fn read_struct_lba<'b, T: Pod + 'b>(
    dev: &dyn BlockDevice,
    lba: u64,
    buffer: &'b mut [u8],
) -> Result<&'b mut T, Error> {
    dev.read_exact(lba * 512, buffer).await?;
    Ok(bytemuck::from_bytes_mut(buffer))
}

pub(crate) async fn probe_gpt(dev: &Arc<dyn BlockDevice>) -> Result<Option<Vec<Partition>>, Error> {
    let mut header = [0; size_of::<GptHeader>()];
    let header = unsafe { read_struct_lba::<GptHeader>(dev.as_ref(), 1, &mut header) }.await?;

    if &header.signature != b"EFI PART" {
        // Not a GPT partition table
        return Ok(None);
    }

    let pt_entsize = header.partition_table_entry_size as usize;
    let pt_len = header.partition_table_len as usize;
    let mut pt_data = PageBox::new_slice(0, pt_len * pt_entsize)?;

    log::info!("pt_entsize = {}, pt_len = {}", pt_entsize, pt_len);

    assert!(size_of::<GptEntry>() <= pt_entsize);

    dev.read_exact(header.partition_table_lba * 512, &mut pt_data)
        .await?;

    let mut partitions = Vec::new();
    for i in 0..pt_len {
        let pt_entry_data = &pt_data[i * pt_entsize..i * pt_entsize + size_of::<GptEntry>()];
        let pt_entry: &GptEntry = bytemuck::from_bytes(pt_entry_data);

        if pt_entry.type_guid.is_nil() {
            break;
        }

        log::debug!(
            "Found partition: guid={}, start={}, end={}",
            pt_entry.type_guid,
            pt_entry.lba_start,
            pt_entry.lba_end
        );
        partitions.push(Partition {
            device: dev.clone(),
            lba_start: pt_entry.lba_start,
            lba_end: pt_entry.lba_end,
        });
    }

    Ok(Some(partitions))
}
