use crate::{drivers, fs, memory::PhysicalMemoryManager};

pub struct StorageInitResult {
    pub ahci_disks: usize,
    pub ext4_mounted: bool,
    pub userspace_image: Option<fs::uefi::LoadedFile>,
    pub userspace_path: Option<&'static [u8]>,
}

pub fn init(memory: &mut PhysicalMemoryManager) -> StorageInitResult {
    let ahci_disks = unsafe { drivers::storage::ahci::init(memory) };
    let mut result = StorageInitResult {
        ahci_disks,
        ext4_mounted: false,
        userspace_image: None,
        userspace_path: None,
    };

    if ahci_disks == 0 {
        return result;
    }

    let Some(mut disk) = drivers::storage::ahci::first_disk() else {
        return result;
    };
    let Ok(partition) = drivers::storage::partition::find_linux_partition(&mut disk) else {
        return result;
    };

    let root = drivers::storage::partition::PartitionBlockDevice::new(disk, partition);
    let root_sector_count = partition.last_lba - partition.first_lba + 1;
    if unsafe { fs::lwext4::mount_root(root, root_sector_count) }.is_err() {
        return result;
    }
    fs::lwext4::init_cache(memory as *mut PhysicalMemoryManager);

    result.ext4_mounted = true;
    let _ = fs::lwext4::mount_vfs_tree();

    match fs::lwext4::load_path(memory, b"/usr/bin/busybox") {
        Ok(image) => {
            result.userspace_image = Some(image);
            result.userspace_path = Some(b"/usr/bin/busybox");
        }
        Err(_) => {
            if let Ok(image) = fs::lwext4::load_path(memory, b"/usr/bin/sh") {
                result.userspace_image = Some(image);
                result.userspace_path = Some(b"/usr/bin/sh");
            }
        }
    }

    result
}