//! Multiboot data structures
//!
//! The [`Header`] is compiled into the kernel and tells Multiboot how to load it.
//!
//! During boot, Multiboot creates a [`Info`] structure for the loaded kernel.
#![allow(unused)]

use core::ffi::CStr;
use core::fmt::{self, Debug};
use core::marker::PhantomData;
use core::{cmp, slice};

use bitfield_struct::bitfield;

use crate::util::RefStr;

/// Multiboot header which tells the bootloader how to load our kernel
///
/// | Offset | Field Name    | Note              |
/// | ------ | ------------- | ----------------- |
/// | 0      | magic         | required          |
/// | 4      | flags         | required          |
/// | 8      | checksum      | required          |
/// | 12     | header_addr   | if flag 16 is set |
/// | 16     | load_addr     | if flag 16 is set |
/// | 20     | load_end_addr | if flag 16 is set |
/// | 24     | bss_end_addr  | if flag 16 is set |
/// | 28     | entry_addr    | if flag 16 is set |
/// | 32     | mode_type     | if flag 2 is set  |
/// | 36     | width         | if flag 2 is set  |
/// | 40     | height        | if flag 2 is set  |
/// | 44     | depth         | if flag 2 is set  |
#[repr(C)]
pub struct Header {
    magic: u32,
    /// 0: page align, 1: mem info, 2: video mode
    flags: u32,
    checksum: u32,
    _unused: [u32; 5],
    /// 0: graphic, 1: text
    video_mode: u32,
    video_width: u32,
    video_height: u32,
    video_depth: u32,
}
impl Header {
    pub const fn new(flags: u32, video_width: u32, video_height: u32, video_depth: u32) -> Self {
        Self {
            magic: Self::MAGIC,
            flags,
            checksum: -(Self::MAGIC as i32 + flags as i32) as u32,
            _unused: [0; 5],
            video_mode: 0,
            video_width,
            video_height,
            video_depth,
        }
    }
}

impl Header {
    pub const MAGIC: u32 = 0x1badb002;
    /// Enforce page alignment for our kernel
    pub const PAGE_ALIGN: u32 = 0b1;
    /// Request info about the physical memory map ([`Info::mem`])
    pub const MEM_INFO: u32 = 0b10;
    /// Enable video mode
    pub const VIDEO_MODE: u32 = 0b100;
}

/// Representation of Multiboot Information according to specification.
#[repr(C)]
pub struct Info {
    pub flags: Flags,

    mem: Memory,
    boot_device: BootDevice,
    cmdline: RefStr,
    mods_count: u32,
    mods_addr: *const Module,
    symbols: SymbolsUnion,

    mmap_length: u32,
    mmap_addr: *const MemoryEntry,
    drives_length: u32,
    drives_addr: u32,

    _config_table: u32,
    boot_loader_name: RefStr,
    _apm_table: u32,

    vbe: VBETable,
    framebuffer: FramebufferTable,
}

#[bitfield(u32)]
pub struct Flags {
    pub memory: bool,
    pub bootdev: bool,
    pub cmdline: bool,
    pub mods: bool,
    pub aout_syms: bool,
    pub elf_shdr: bool,
    pub memory_map: bool,
    pub drive_info: bool,
    pub config_table: bool,
    pub boot_loader_name: bool,
    pub apm_table: bool,
    pub vbe: bool,
    pub framebuffer: bool,
    #[bits(19)]
    __: (),
}

impl Info {
    /// Parses the Multiboot information structures
    pub unsafe fn from_ptr(ptr: *const u8) -> &'static Self {
        &*ptr.cast()
    }
    /// Returns the memory sizes
    pub fn memory(&self) -> Option<&Memory> {
        self.flags.memory().then_some(&self.mem)
    }
    /// Indicates which bios disk device the boot loader loaded the OS image from.
    pub fn boot_device(&self) -> Option<&BootDevice> {
        self.flags.bootdev().then_some(&self.boot_device)
    }
    /// Command line to be passed to the kernel.
    pub fn cmdline(&self) -> Option<&str> {
        self.flags.cmdline().then_some(self.cmdline.as_str())
    }
    /// Get the name of the bootloader.
    pub fn boot_loader_name(&self) -> Option<&str> {
        self.flags
            .boot_loader_name()
            .then_some(self.boot_loader_name.as_str())
    }
    /// Discover all additional modules in multiboot.
    pub fn modules(&self) -> Option<&[Module]> {
        if self.flags.mods() {
            Some(unsafe { slice::from_raw_parts(self.mods_addr as _, self.mods_count as _) })
        } else {
            None
        }
    }
    /// Get the symbols.
    pub fn symbols(&self) -> Option<Symbols> {
        unsafe {
            match (self.flags.aout_syms(), self.flags.elf_shdr()) {
                (true, false) => Some(Symbols::Elf(self.symbols.elf)),
                (false, true) => Some(Symbols::AOut(self.symbols.aout)),
                _ => None,
            }
        }
    }
    /// Discover all memory regions in the multiboot memory map.
    pub fn memory_map(&self) -> Option<MemoryMapIter> {
        if self.flags.memory_map() {
            Some(MemoryMapIter {
                addr: self.mmap_addr,
                end: (self.mmap_addr as u32 + self.mmap_length) as _,
                __: PhantomData,
            })
        } else {
            None
        }
    }
    /// Return end address of multiboot image.
    ///
    /// This function can be used to figure out a (hopefully) safe offset
    /// in the first region of memory to start using as free memory.
    pub fn find_highest_address(&self) -> u32 {
        self.cmdline()
            .map_or(0, |f| f.as_ptr() as u32 + f.len() as u32)
            .max(
                self.boot_loader_name()
                    .map_or(0, |f| f.as_ptr() as u32 + f.len() as u32),
            )
            .max(match self.symbols() {
                Some(Symbols::Elf(e)) => e.addr + e.num * e.size,
                Some(Symbols::AOut(a)) => {
                    a.addr + a.tabsize + a.strsize + 2 * core::mem::size_of::<u32>() as u32
                }
                None => 0,
            })
            .max(self.mmap_addr as u32 + self.mmap_length)
            .max(self.drives_addr + self.drives_length)
            .max(self.mods_addr as u32 + self.mods_count * core::mem::size_of::<Module>() as u32)
            .max(
                self.modules()
                    .into_iter()
                    .flatten()
                    .map(|m| m.end)
                    .max()
                    .unwrap_or(0),
            )
            .next_multiple_of(4096)
    }
    /// Contains the VESA BIOS extensions
    pub fn vbe(&self) -> Option<&VBETable> {
        self.flags.vbe().then_some(&self.vbe)
    }
    /// Contains the information about the framebuffer
    pub fn framebuffer(&self) -> Option<&FramebufferTable> {
        self.flags.framebuffer().then_some(&self.framebuffer)
    }
}

impl fmt::Debug for Info {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Information")
            .field("flags", &self.flags)
            .field("mem", &self.memory())
            .field("boot_device", &self.boot_device())
            .field("cmdline", &self.cmdline())
            .field("mods_count", &self.mods_count)
            .field("mods_addr", &self.mods_addr)
            .field("symbols", &self.symbols())
            .field("mmap_length", &self.mmap_length)
            .field("mmap_addr", &self.mmap_addr)
            .field("drives_length", &self.drives_length)
            .field("drives_addr", &self.drives_addr)
            .field("config_table", &self._config_table)
            .field("boot_loader_name", &self.boot_loader_name())
            .field("apm_table", &self._apm_table)
            .field("vbe", &self.vbe())
            .field("framebuffer", &self.framebuffer())
            .finish()
    }
}

/// ‘mem_lower’ and ‘mem_upper’ indicate the amount of lower and upper memory, respectively, in kilobytes.
#[repr(C)]
#[derive(Debug)]
pub struct Memory {
    /// Lower memory starts at address 0, and upper memory starts at address 1 megabyte. The maximum possible value for lower memory is 640 kilobytes.
    pub lower: u32,
    /// The value returned for upper memory is maximally the address of the first upper memory hole minus 1 megabyte. It is not guaranteed to be this value.
    pub upper: u32,
}

#[derive(Debug, Clone)]
#[repr(C)]
pub struct BootDevice {
    /// Contains the bios drive number as understood by
    /// the bios INT 0x13 low-level disk interface: e.g. 0x00 for the
    /// first floppy disk or 0x80 for the first hard disk.
    pub drive: u8,
    /// Specifies the top-level partition number.
    pub partition1: u8,
    /// Specifies a sub-partition in the top-level partition
    pub partition2: u8,
    /// Specifies a sub-partition in the 2nd-level partition
    pub partition3: u8,
}

/// Multiboot format to information about module
#[repr(C)]
#[derive(Debug)]
pub struct Module {
    /// Start address of module in memory.
    pub start: u32,

    /// End address of module in memory.
    pub end: u32,

    /// The `string` field provides an arbitrary string to be associated
    /// with that particular boot module.
    ///
    /// It is a zero-terminated ASCII string, just like the kernel command line.
    /// The `string` field may be 0 if there is no string associated with the module.
    /// Typically the string might be a command line (e.g. if the operating system
    /// treats boot modules as executable programs), or a pathname
    /// (e.g. if the operating system treats boot modules as files in a file system),
    /// but its exact use is specific to the operating system.
    pub string: RefStr,

    /// Must be zero.
    reserved: u32,
}

/// Multiboot format for Symbols
#[repr(C)]
union SymbolsUnion {
    aout: AOutSymbols,
    elf: ElfSymbols,
    _align: [u32; 4usize],
}
#[derive(Debug)]
pub enum Symbols {
    AOut(AOutSymbols),
    Elf(ElfSymbols),
}

/// Multiboot format for AOut Symbols
#[repr(C)]
#[derive(Default, Copy, Clone, Debug)]
pub struct AOutSymbols {
    pub tabsize: u32,
    pub strsize: u32,
    pub addr: u32,
    pub reserved: u32,
}

/// Multiboot format for ELF Symbols
#[repr(C)]
#[derive(Default, Copy, Clone, Debug)]
pub struct ElfSymbols {
    pub num: u32,
    pub size: u32,
    pub addr: u32,
    pub shndx: u32,
}

/// Types that define if the memory is usable or not.
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
#[allow(clippy::upper_case_acronyms)]
#[repr(u32)]
pub enum MemoryType {
    /// memory, available to OS
    Available = 1,
    /// reserved, not available (rom, mem map dev)
    Reserved = 2,
    /// ACPI Reclaim Memory
    ACPI = 3,
    /// ACPI NVS Memory
    NVS = 4,
    /// defective RAM modules
    Defect = 5,
}

/// Multiboot format of the MMAP buffer.
///
/// Note that size is defined to be at -4 bytes in multiboot.
#[derive(Debug, Clone)]
#[repr(C, packed)]
pub struct MemoryEntry {
    size: u32,
    pub base_addr: u64,
    pub length: u64,
    pub ty: MemoryType,
}

pub struct MemoryMapIter<'a> {
    addr: *const MemoryEntry,
    end: *const MemoryEntry,
    __: PhantomData<&'a ()>,
}
impl<'a> Iterator for MemoryMapIter<'a> {
    type Item = &'a MemoryEntry;

    fn next(&mut self) -> Option<Self::Item> {
        if self.addr < self.end {
            assert!(!self.addr.is_null());
            let entry = unsafe { &*(self.addr as *const MemoryEntry) };
            assert!(entry.size >= 20);
            self.addr = (self.addr as u32 + entry.size + size_of::<u32>() as u32) as _;
            Some(entry)
        } else {
            None
        }
    }
}

/// Contains information about the VESA BIOS extension
#[derive(Debug)]
pub struct VBETable {
    pub control_info: u32,
    pub mode_info: u32,
    pub mode: u16,
    pub interface_seg: u16,
    pub interface_off: u16,
    pub interface_len: u16,
}

/// Contains the information about the framebuffer
#[derive(Clone)]
#[repr(C)]
pub struct FramebufferTable {
    pub addr: u64,
    pub pitch: u32,
    pub width: u32,
    pub height: u32,
    pub bpp: u8,
    ty: u8,
    color_info: ColorInfoUnion,
}

impl FramebufferTable {
    pub fn new(
        addr: u64,
        pitch: u32,
        width: u32,
        height: u32,
        bpp: u8,
        color_info: Option<ColorInfo>,
    ) -> Self {
        let (ty, color_info) = match color_info {
            Some(ColorInfo::Palette(palette)) => (0, ColorInfoUnion { palette }),
            Some(ColorInfo::Rgb(rgb)) => (1, ColorInfoUnion { rgb }),
            Some(ColorInfo::Text) => (2, ColorInfoUnion { _align: [0; 2] }),
            None => todo!(),
        };
        Self {
            addr,
            pitch,
            width,
            height,
            bpp,
            ty,
            color_info,
        }
    }
    pub fn color_info(&self) -> Option<ColorInfo> {
        unsafe {
            match self.ty {
                0 => Some(ColorInfo::Palette(self.color_info.palette)),
                1 => Some(ColorInfo::Rgb(self.color_info.rgb)),
                2 => Some(ColorInfo::Text),
                _ => None,
            }
        }
    }
}

impl fmt::Debug for FramebufferTable {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("FramebufferTable")
            .field("addr", &self.addr)
            .field("pitch", &self.pitch)
            .field("width", &self.width)
            .field("height", &self.height)
            .field("bpp", &self.bpp)
            .field("color_info", &self.color_info())
            .finish()
    }
}

#[derive(Debug)]
pub enum ColorInfo {
    Palette(ColorInfoPalette),
    Rgb(ColorInfoRgb),
    Text,
}

/// Multiboot format for the framebuffer color info
#[repr(C)]
#[derive(Clone, Copy)]
union ColorInfoUnion {
    palette: ColorInfoPalette,
    rgb: ColorInfoRgb,
    _align: [u32; 2usize],
}

/// Information for indexed color mode
#[repr(C)]
#[derive(Debug, Clone, Copy)]
pub struct ColorInfoPalette {
    pub palette_addr: u32,
    pub palette_num_colors: u16,
}

/// Information for direct RGB color mode
#[repr(C)]
#[derive(Debug, Clone, Copy)]
pub struct ColorInfoRgb {
    pub red_offset: u8,
    pub red_bits: u8,
    pub green_offset: u8,
    pub green_bits: u8,
    pub blue_offset: u8,
    pub blue_bits: u8,
}