//! Partition-table layer.
//!
//! Two implementations in v1: [`mbr::Mbr`] (4-primary MBR) and [`gpt::Gpt`]
//! (128-entry GPT with protective MBR + backup header). Both go through the
//! [`PartitionTable`] trait so higher layers can target either uniformly.
//!
//! A partition table is *just metadata* — it does not allocate or zero data
//! blocks beyond its own header sectors. To get a filesystem-grade view of a
//! single partition's bytes, call [`slice_partition`] which builds a
//! [`SlicedBackend`] covering exactly the partition's LBA range.

use uuid::Uuid;

use crate::Result;
use crate::block::{BlockDevice, SlicedBackend};

pub mod gpt;
pub mod mbr;

pub use gpt::Gpt;
pub use mbr::Mbr;

/// One partition entry, unified across MBR and GPT.
///
/// MBR ignores `uuid`, `name`, and `attributes`; the high bits of `start_lba`
/// and `size_lba` must fit in 32 bits or [`Mbr::write`] returns
/// [`crate::Error::Unsupported`]. GPT writes them all.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Partition {
    /// First LBA of the partition (inclusive).
    pub start_lba: u64,
    /// Length of the partition, in LBAs.
    pub size_lba: u64,
    /// Semantic / on-disk type.
    pub kind: PartitionKind,
    /// Per-partition UUID. Auto-generated by [`Gpt::build`] when `None`; MBR
    /// ignores it.
    pub uuid: Option<Uuid>,
    /// Label, up to 36 UTF-16 code units after encoding. GPT-only.
    pub name: Option<String>,
    /// MBR active flag (sets entry boot byte to 0x80). On GPT the writer maps
    /// it onto attribute bit 2 ("legacy BIOS bootable").
    pub bootable: bool,
    /// Raw 64-bit attribute bitfield. GPT-only. Bit 2 is OR'd with `bootable`.
    pub attributes: u64,
}

impl Partition {
    /// Partition that occupies LBAs `[start, start+size)`, of `kind`, with all
    /// other fields defaulted.
    pub fn new(start_lba: u64, size_lba: u64, kind: PartitionKind) -> Self {
        Self {
            start_lba,
            size_lba,
            kind,
            uuid: None,
            name: None,
            bootable: false,
            attributes: 0,
        }
    }

    /// Last LBA of the partition (inclusive). Used heavily by GPT, which
    /// stores ending LBAs rather than sizes.
    pub fn end_lba(&self) -> u64 {
        self.start_lba + self.size_lba - 1
    }
}

/// Semantic partition type. Convertible to either an MBR `u8` code or a GPT
/// type UUID; an escape-hatch variant exists for each direction.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum PartitionKind {
    /// Empty / unused entry. MBR 0x00, GPT all-zero UUID.
    Empty,
    /// Generic Linux filesystem (ext, xfs, btrfs, ...). MBR 0x83.
    LinuxFilesystem,
    /// Linux swap. MBR 0x82.
    LinuxSwap,
    /// EFI System Partition (FAT32 by convention).
    EfiSystem,
    /// BIOS boot partition (GRUB stage-1.5 storage for BIOS-on-GPT). GPT only.
    BiosBoot,
    /// FAT32 with LBA addressing. MBR 0x0C.
    Fat32,
    /// Microsoft Basic Data — generic FAT/NTFS on GPT.
    MicrosoftBasicData,
    /// Raw MBR type code (escape hatch for unmapped types).
    Mbr(u8),
    /// Raw GPT type UUID (escape hatch for unmapped types).
    Gpt(Uuid),
}

impl PartitionKind {
    /// Map to the MBR partition-type byte. For variants that have no MBR
    /// equivalent (e.g. `BiosBoot`), returns `0x00`.
    pub fn as_mbr_byte(&self) -> u8 {
        match self {
            PartitionKind::Empty => 0x00,
            PartitionKind::LinuxFilesystem => 0x83,
            PartitionKind::LinuxSwap => 0x82,
            PartitionKind::EfiSystem => 0xEF,
            PartitionKind::Fat32 => 0x0C,
            PartitionKind::MicrosoftBasicData => 0x07,
            PartitionKind::Mbr(b) => *b,
            // No MBR mapping; caller should treat 0 as "leave empty" and pick
            // a different kind explicitly if writing MBR.
            PartitionKind::BiosBoot => 0x00,
            PartitionKind::Gpt(_) => 0x00,
        }
    }

    /// Build from an MBR partition-type byte. Unknown bytes round-trip as
    /// `Mbr(b)`.
    pub fn from_mbr_byte(b: u8) -> Self {
        match b {
            0x00 => PartitionKind::Empty,
            0x83 => PartitionKind::LinuxFilesystem,
            0x82 => PartitionKind::LinuxSwap,
            0xEF => PartitionKind::EfiSystem,
            0x0C => PartitionKind::Fat32,
            0x07 => PartitionKind::MicrosoftBasicData,
            other => PartitionKind::Mbr(other),
        }
    }

    /// Map to the GPT type UUID. Returns the nil UUID for variants that have
    /// no GPT mapping (e.g. raw MBR codes).
    pub fn as_gpt_uuid(&self) -> Uuid {
        match self {
            PartitionKind::Empty => Uuid::nil(),
            // Standard GPT type UUIDs per UEFI spec / discoverable-partitions
            // spec.
            PartitionKind::LinuxFilesystem => uuid_const("0FC63DAF-8483-4772-8E79-3D69D8477DE4"),
            PartitionKind::LinuxSwap => uuid_const("0657FD6D-A4AB-43C4-84E5-0933C84B4F4F"),
            PartitionKind::EfiSystem => uuid_const("C12A7328-F81F-11D2-BA4B-00A0C93EC93B"),
            PartitionKind::BiosBoot => uuid_const("21686148-6449-6E6F-744E-656564454649"),
            PartitionKind::Fat32 | PartitionKind::MicrosoftBasicData => {
                uuid_const("EBD0A0A2-B9E5-4433-87C0-68B6B72699C7")
            }
            PartitionKind::Gpt(u) => *u,
            PartitionKind::Mbr(_) => Uuid::nil(),
        }
    }

    /// Build from a GPT type UUID. Unknown UUIDs round-trip as `Gpt(uuid)`.
    pub fn from_gpt_uuid(u: Uuid) -> Self {
        if u.is_nil() {
            return PartitionKind::Empty;
        }
        match u.as_hyphenated().to_string().to_ascii_uppercase().as_str() {
            "0FC63DAF-8483-4772-8E79-3D69D8477DE4" => PartitionKind::LinuxFilesystem,
            "0657FD6D-A4AB-43C4-84E5-0933C84B4F4F" => PartitionKind::LinuxSwap,
            "C12A7328-F81F-11D2-BA4B-00A0C93EC93B" => PartitionKind::EfiSystem,
            "21686148-6449-6E6F-744E-656564454649" => PartitionKind::BiosBoot,
            "EBD0A0A2-B9E5-4433-87C0-68B6B72699C7" => PartitionKind::MicrosoftBasicData,
            _ => PartitionKind::Gpt(u),
        }
    }
}

/// Common interface implemented by [`Mbr`] and [`Gpt`]. Object-safe enough
/// that higher layers can hold a `Box<dyn PartitionTable>` if useful.
pub trait PartitionTable {
    /// Write the partition table to its conventional location on `dev`.
    /// `dev.total_size()` must be at least large enough to hold the table's
    /// metadata sectors (1 for MBR; LBA 0..33 plus the trailing backup
    /// header for GPT).
    fn write(&self, dev: &mut dyn BlockDevice) -> Result<()>;

    /// All non-empty partition entries, in disk order.
    fn partitions(&self) -> &[Partition];
}

/// Build a [`SlicedBackend`] that exposes exactly the bytes of `pt.partitions()[index]`.
///
/// Returned device is sized in bytes: `partition.size_lba * dev.block_size()`.
pub fn slice_partition<'a, B: BlockDevice + ?Sized, P: PartitionTable + ?Sized>(
    pt: &P,
    dev: &'a mut B,
    index: usize,
) -> Result<SlicedBackend<'a, B>> {
    let parts = pt.partitions();
    let p = parts.get(index).ok_or_else(|| {
        crate::Error::InvalidArgument(format!(
            "partition index {index} out of range (have {})",
            parts.len()
        ))
    })?;
    let bs = u64::from(dev.block_size());
    SlicedBackend::new(dev, p.start_lba * bs, p.size_lba * bs)
}

/// Parse a UUID from a `const` hyphenated string. Panics on malformed input
/// — used only with hard-coded known-good GPT type UUIDs.
fn uuid_const(s: &str) -> Uuid {
    Uuid::parse_str(s).expect("hard-coded UUID literal")
}

/// Encode a UUID into the 16-byte on-disk GPT representation: time-low,
/// time-mid, and time-hi-and-version are little-endian; the remaining 8 bytes
/// are written as-is.
pub(crate) fn uuid_to_gpt_bytes(u: Uuid) -> [u8; 16] {
    let f = u.as_fields();
    let mut out = [0u8; 16];
    out[0..4].copy_from_slice(&f.0.to_le_bytes());
    out[4..6].copy_from_slice(&f.1.to_le_bytes());
    out[6..8].copy_from_slice(&f.2.to_le_bytes());
    out[8..16].copy_from_slice(f.3);
    out
}

/// Decode a 16-byte GPT on-disk UUID into a [`Uuid`].
pub(crate) fn uuid_from_gpt_bytes(b: &[u8; 16]) -> Uuid {
    let d1 = u32::from_le_bytes(b[0..4].try_into().unwrap());
    let d2 = u16::from_le_bytes(b[4..6].try_into().unwrap());
    let d3 = u16::from_le_bytes(b[6..8].try_into().unwrap());
    let mut d4 = [0u8; 8];
    d4.copy_from_slice(&b[8..16]);
    Uuid::from_fields(d1, d2, d3, &d4)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn uuid_endian_roundtrip() {
        let cases = [
            "C12A7328-F81F-11D2-BA4B-00A0C93EC93B",
            "0FC63DAF-8483-4772-8E79-3D69D8477DE4",
            "21686148-6449-6E6F-744E-656564454649",
        ];
        for s in cases {
            let u = Uuid::parse_str(s).unwrap();
            let bytes = uuid_to_gpt_bytes(u);
            let back = uuid_from_gpt_bytes(&bytes);
            assert_eq!(u, back, "roundtrip for {s}");
        }
    }

    #[test]
    fn efi_uuid_disk_layout_is_correct() {
        // From the UEFI spec, the ESP UUID C12A7328-F81F-11D2-BA4B-00A0C93EC93B
        // serialises on disk as: 28 73 2A C1 1F F8 D2 11 BA 4B 00 A0 C9 3E C9 3B
        let u = Uuid::parse_str("C12A7328-F81F-11D2-BA4B-00A0C93EC93B").unwrap();
        let bytes = uuid_to_gpt_bytes(u);
        let expected = [
            0x28, 0x73, 0x2A, 0xC1, 0x1F, 0xF8, 0xD2, 0x11, 0xBA, 0x4B, 0x00, 0xA0, 0xC9, 0x3E,
            0xC9, 0x3B,
        ];
        assert_eq!(bytes, expected);
    }

    #[test]
    fn mbr_byte_roundtrip() {
        for b in 0u8..=255 {
            let k = PartitionKind::from_mbr_byte(b);
            // Known mapped types should re-encode to the same byte.
            // Unknown ones round-trip via Mbr(b).
            assert_eq!(k.as_mbr_byte(), b, "byte {b:#04x} round-trip");
        }
    }
}
