//! Partition tables: MBR (1983, max 4 primary partitions, 2 TiB) and GPT
//! (UEFI era, 128 partitions, 64-bit sizes). A [`Partition`] is itself a
//! `BlockDevice` -- a window onto part of the disk -- so a file system on
//! "vda1" doesn't know (or care) that it starts at sector 2048 of "vda".
//!
//! MBR (sector 0): boot code, then 4 entries of 16 bytes at offset 446, then
//! the signature 0x55 0xAA. Each entry: type byte (offset 4), first sector
//! (u32 at 8), sector count (u32 at 12). Type 0xEE = "this disk is GPT, the
//! MBR is just a protective placeholder".
//!
//! GPT: header in sector 1 ("EFI PART"; where the entry array is, how many
//! entries, how big each is), entry array usually in sectors 2-33. Each
//! entry: type GUID (all zero = unused), unique GUID, first/last LBA, flags,
//! UTF-16 name. A backup copy lives at the end of the disk; we don't check
//! it, nor the CRC32 checksums (exercise!). Extended/logical MBR partitions
//! aren't supported either. References: wiki.osdev.org/MBR_(x86),
//! wiki.osdev.org/GPT

use super::{BlockDevice, SECTOR_SIZE, check_range};
use crate::errno::Errno;
use alloc::{string::String, sync::Arc, vec, vec::Vec};

/// What the partition table says about one partition.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PartitionInfo {
    pub start: u64,   // first sector
    pub sectors: u64, // length
    pub type_name: &'static str,
    pub name: String, // GPT only
}

#[derive(Debug, PartialEq, Eq)]
pub enum PartitionError {
    NoSignature, // sector 0 doesn't end in 55 AA: not partitioned
    /// Sector 0 ends in 55 AA but is an exFAT boot sector, not an MBR: the
    /// whole disk is one exFAT volume.
    WholeDiskExfat,
    BadGpt,
    Io(Errno),
}

fn u32_at(b: &[u8], at: usize) -> u32 {
    u32::from_le_bytes(b[at..at + 4].try_into().unwrap())
}
fn u64_at(b: &[u8], at: usize) -> u64 {
    u64::from_le_bytes(b[at..at + 8].try_into().unwrap())
}

/// Read the partition table of `disk`: ("MBR" or "GPT", partitions).
pub fn scan(disk: &Arc<dyn BlockDevice>) -> Result<(&'static str, Vec<PartitionInfo>), PartitionError> {
    let mut mbr = [0u8; SECTOR_SIZE];
    disk.read_sectors(0, &mut mbr).map_err(PartitionError::Io)?;
    if mbr[510..512] != [0x55, 0xaa] {
        return Err(PartitionError::NoSignature);
    }
    // A file system's boot sector also ends in 55 AA. exFAT's is easy to
    // tell apart by its name at byte 3 (and its partition table area is
    // just zeros, which would read as "MBR, 0 partitions").
    if &mbr[3..11] == b"EXFAT   " {
        return Err(PartitionError::WholeDiskExfat);
    }
    let entries: Vec<&[u8]> = (0..4).map(|i| &mbr[446 + 16 * i..446 + 16 * (i + 1)]).collect();
    if entries.iter().any(|e| e[4] == 0xee) {
        return Ok(("GPT", scan_gpt(disk)?));
    }
    let parts = entries
        .iter()
        .filter(|e| e[4] != 0 && u32_at(e, 12) != 0) // type 0 = empty slot
        .map(|e| PartitionInfo {
            start: u32_at(e, 8) as u64,
            sectors: u32_at(e, 12) as u64,
            type_name: mbr_type_name(e[4]),
            name: String::new(),
        })
        .collect();
    Ok(("MBR", parts))
}

fn scan_gpt(disk: &Arc<dyn BlockDevice>) -> Result<Vec<PartitionInfo>, PartitionError> {
    let mut header = [0u8; SECTOR_SIZE];
    disk.read_sectors(1, &mut header).map_err(PartitionError::Io)?;
    if &header[0..8] != b"EFI PART" {
        return Err(PartitionError::BadGpt);
    }
    let entries_lba = u64_at(&header, 72);
    let count = u32_at(&header, 80) as usize;
    let entry_size = u32_at(&header, 84) as usize;
    if entry_size < 128 || count > 1024 || entry_size > SECTOR_SIZE || SECTOR_SIZE % entry_size != 0 {
        return Err(PartitionError::BadGpt);
    }
    let sectors = (count * entry_size).div_ceil(SECTOR_SIZE);
    let mut table = vec![0u8; sectors * SECTOR_SIZE];
    disk.read_sectors(entries_lba, &mut table).map_err(PartitionError::Io)?;

    let mut parts = Vec::new();
    for e in table.chunks(entry_size).take(count) {
        let type_guid = &e[0..16];
        if type_guid.iter().all(|&b| b == 0) {
            continue; // unused entry
        }
        let (first, last) = (u64_at(e, 32), u64_at(e, 40));
        if last < first {
            return Err(PartitionError::BadGpt);
        }
        // Name: UTF-16LE, up to 36 code units, zero-padded.
        let units: Vec<u16> = e[56..128].chunks(2).map(|c| u16::from_le_bytes([c[0], c[1]])).take_while(|&u| u != 0).collect();
        parts.push(PartitionInfo {
            start: first,
            sectors: last - first + 1,
            type_name: gpt_type_name(type_guid),
            name: char::decode_utf16(units).map(|c| c.unwrap_or('?')).collect(),
        });
    }
    Ok(parts)
}

fn mbr_type_name(t: u8) -> &'static str {
    match t {
        0x0b | 0x0c => "FAT32",
        0x07 => "NTFS/exFAT",
        0x83 => "Linux",
        0x05 | 0x0f => "extended (unsupported)",
        _ => "unknown",
    }
}

/// GUIDs are stored with their first three fields little-endian, so the
/// bytes look scrambled compared to the usual text form.
fn gpt_type_name(guid: &[u8]) -> &'static str {
    const BASIC_DATA: [u8; 16] = [0xa2, 0xa0, 0xd0, 0xeb, 0xe5, 0xb9, 0x33, 0x44, 0x87, 0xc0, 0x68, 0xb6, 0xb7, 0x26, 0x99, 0xc7]; // EBD0A0A2-B9E5-4433-87C0-68B6B72699C7
    const LINUX_FS: [u8; 16] = [0xaf, 0x3d, 0xc6, 0x0f, 0x83, 0x84, 0x72, 0x47, 0x8e, 0x79, 0x3d, 0x69, 0xd8, 0x47, 0x7d, 0xe4]; // 0FC63DAF-8483-4772-8E79-3D69D8477DE4
    const EFI_SYSTEM: [u8; 16] = [0x28, 0x73, 0x2a, 0xc1, 0x1f, 0xf8, 0xd2, 0x11, 0xba, 0x4b, 0x00, 0xa0, 0xc9, 0x3e, 0xc9, 0x3b]; // C12A7328-F81F-11D2-BA4B-00A0C93EC93B
    match guid {
        g if g == BASIC_DATA => "Microsoft basic data (FAT/exFAT/NTFS)",
        g if g == LINUX_FS => "Linux filesystem",
        g if g == EFI_SYSTEM => "EFI system (FAT)",
        _ => "unknown",
    }
}

/// A partition as a block device: sector numbers are shifted by `start` and
/// checked against `sectors`, so a file system can't scribble outside it.
pub struct Partition {
    disk: Arc<dyn BlockDevice>,
    start: u64,
    sectors: u64,
}

impl Partition {
    pub fn new(disk: Arc<dyn BlockDevice>, start: u64, sectors: u64) -> Self {
        Partition { disk, start, sectors }
    }
}

impl BlockDevice for Partition {
    fn sector_count(&self) -> u64 {
        self.sectors
    }
    fn read_sectors(&self, lba: u64, buf: &mut [u8]) -> Result<(), Errno> {
        check_range(self, lba, buf.len())?;
        self.disk.read_sectors(self.start + lba, buf)
    }
    fn write_sectors(&self, lba: u64, buf: &[u8]) -> Result<(), Errno> {
        check_range(self, lba, buf.len())?;
        self.disk.write_sectors(self.start + lba, buf)
    }
}