//! A Label Storage Area

mod label_index_block;
pub mod namespace_label;
pub mod region_label;

use crate::lsa::{
    label_index_block::LabelIndexBlock,
    namespace_label::*,
    region_label::*,
    Label::{Empty, Namespace, Region, Vendor},
};
use serde::{
    ser::{SerializeStruct, Serializer},
    Serialize,
};
use std::{
    fmt,
    io::{Error, ErrorKind, Write},
    slice,
};
use uuid::Uuid;

/// Sets the fletcher64 checksum for a label that has a field named `Checksum`.
#[macro_export]
macro_rules! set_checksum {
    ( $self:ident ) => {
        $self.Checksum = 0;
        let p: *const Self = $self;
        let p: *const u32 = p as *const u32;
        let s: &[u32] = unsafe { slice::from_raw_parts(p, 4) };
        $self.Checksum = generate_checksum(s);
    };
}

///Verify a checksum for a label that has a field named `Checksum`
#[macro_export]
macro_rules! verify_checksum {
    ( $self:ident ) => {{
        let mut temp = *$self;
        temp.Checksum = 0;
        let p: *const Self = &temp;
        let p: *const u32 = p as *const _;
        let dwords = unsafe { slice::from_raw_parts(p, 4) };
        if $self.Checksum != generate_checksum(dwords) {
            Err(Error::new(ErrorKind::Other, "Invalid checksum"))
        } else {
            Ok(())
        }
    }};
}

/// Helper macro to turn the packed labels into bytes
#[macro_export]
macro_rules! label_as_bytes {
    ( $self:ident ) => {{
        let p: *const Self = $self;
        let p: *const u8 = p as *const _;
        let bytes: &[u8] = unsafe { slice::from_raw_parts(p, 256) };
        &bytes[..256]
    }};
}

/// Helper serializer for UUIDs
fn to_uuid<S>(x: &[u8; 16], s: S) -> Result<S::Ok, S::Error>
where
    S: Serializer,
{
    s.serialize_str(&Uuid::from_bytes(*x).to_hyphenated().to_string())
}

/// Helper serializer for values which should be 0x...
fn to_hex<S>(x: &u64, s: S) -> Result<S::Ok, S::Error>
where
    S: Serializer,
{
    let hex = format!("0x{:x}", x);
    s.serialize_str(&hex)
}

/// Helper function to get the fletcher64 checksum
fn generate_checksum(dwords: &[u32]) -> u64 {
    let mut checksum = fletcher::Fletcher64::new();
    checksum.update(dwords);
    checksum.value()
}

pub fn label_from_bytes(raw_label: &[u8]) -> Result<Label, std::io::Error> {
    assert!(raw_label.len() == 256);
    let uuid: Uuid = match Uuid::from_slice(&raw_label[..16]) {
        Ok(u) => u,
        Err(e) => return Err(Error::new(ErrorKind::Other, e)),
    };

    // As labels are a fixed size, these should all be safe.
    let label: Label = match uuid {
        REGION_UUID => Region(RegionLabel::from_bytes(raw_label)),
        NAMESPACE_UUID => Namespace(NamespaceLabel::from_bytes(raw_label)),
        _ => Empty,
    };

    Ok(label)
}

/// Properties common to all types of labels
pub trait CommonLabelProperties {
    /// Updates the checksum for the given label
    ///
    /// For anything with a "Checksum" field and using fletcher64, see `set_checksum!()`
    fn update_checksum(&mut self);

    /// Verifies the checksum is correct.
    ///
    /// For anything with a "Checksum" field and using fletcher64, see `verify_checksum!()`
    fn verify(&self) -> Result<(), std::io::Error>;

    /// Exports a labels as a byte slice.
    ///
    /// For standard labels, see `label_as_bytes!()`
    fn as_bytes(&self) -> &[u8];

    /// Imports a label from a byte slice.
    fn from_bytes(buf: &[u8]) -> Self;

    /// Gets the size of the label.
    ///
    /// As of CXL 2.0 spec, all labels must be 256 bytes.
    fn get_size(&self) -> usize {
        256
    }

    fn set_slot(&mut self, size: usize);
}

/// Vendor labels in the Label Storage area.
///
/// Vendor label's data is entirely opaque.
#[repr(packed)]
#[derive(Clone, Copy, Serialize)]
#[allow(non_snake_case, dead_code)]
pub struct VendorLabel {
    #[serde(serialize_with = "to_uuid")]
    Type: [u8; 16],
    #[serde(skip_serializing)]
    Vendor: [u8; 0xe8],
    Checksum: u64,
}

impl CommonLabelProperties for VendorLabel {
    fn update_checksum(&mut self) {
        set_checksum!(self);
    }
    fn verify(&self) -> Result<(), std::io::Error> {
        verify_checksum!(self)
    }
    fn as_bytes(&self) -> &[u8] {
        label_as_bytes!(self)
    }
    fn from_bytes(buf: &[u8]) -> Self {
        assert!(buf.len() == 256);
        unsafe { std::ptr::read(buf.as_ptr() as *const _) }
    }
    fn set_slot(&mut self, _slot: usize) {
        panic!("Invalid set_slot");
    }
}

/// A Label occupying a slot in the Label Storage Area
#[derive(Serialize)]
pub enum Label {
    Region(RegionLabel),
    Namespace(NamespaceLabel),
    Vendor(VendorLabel),
    Empty,
}

impl fmt::Display for Label {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Label::Region(_) => write!(f, "Region Label"),
            Label::Namespace(_) => write!(f, "Namespace Label"),
            Label::Vendor(_) => write!(f, "Vendor Label"),
            Label::Empty => write!(f, "?"),
        }
    }
}

impl CommonLabelProperties for Label {
    fn update_checksum(&mut self) {
        match self {
            Region(r) => r.update_checksum(),
            Namespace(n) => n.update_checksum(),
            Vendor(v) => v.update_checksum(),
            _ => {}
        };
    }
    fn verify(&self) -> Result<(), std::io::Error> {
        match self {
            Region(r) => r.verify(),
            Namespace(n) => n.verify(),
            Vendor(v) => v.verify(),
            _ => Ok(()),
        }
    }
    fn as_bytes(&self) -> &[u8] {
        match self {
            Region(r) => r.as_bytes(),
            Namespace(n) => n.as_bytes(),
            Vendor(v) => v.as_bytes(),
            _ => unreachable!(),
        }
    }

    fn from_bytes(buf: &[u8]) -> Self {
        label_from_bytes(buf).expect("Couldn't decipher label")
    }

    fn set_slot(&mut self, slot: usize) {
        match self {
            Region(r) => r.set_slot(slot),
            Namespace(n) => n.set_slot(slot),
            _ => unreachable!(),
        }
    }
}

fn next_seq(seq: u32) -> u32 {
    assert!(seq < 4);
    let mut next = seq + 1 % 4;
    if next == 0 {
        next = 1;
    }
    next
}

/// Label Storage Area
///
/// A Label Storage area is defined by the CXL 2.0 specification. It is comprised of 2 Label Index
/// Blocks and a set of at least 3 labels.
///
/// "The LSA consists of two Label Index Blocks followed by an array of label slots."
pub struct LabelStorageArea {
    /// Each LSA has 2 Label Index Blocks.
    label_index_blocks: [LabelIndexBlock; 2],
    /// An LSA is comprised of at least 3 label slots
    label_slots: Vec<Label>,

    /// The active Label Index Block.
    ///
    /// Label block indices are ping-ponged back and forth for atomicity. Typically, the inactive
    /// label block should be updated while the active label block should only be read.
    active_lib: isize,

    /// The current "highest" sequence number of the label block indices. When a Label Storage Area
    /// is first initialized, this value will be invalid.
    current_seq: u32,
}

impl LabelStorageArea {
    /// Returns a label storage area that can hold the given number of slots.
    ///
    /// The returned label storage area will be invalid.
    ///
    /// # Arguments
    ///
    /// * `slots`: - The max number of slots to allocate for.
    ///
    /// # Examples
    ///
    /// ```
    /// use cxl_rs::lsa::LabelStorageArea;
    /// let lsa = LabelStorageArea::new(20);
    /// // Create some labels!
    /// ```
    ///
    /// # TODO
    ///
    /// * `interleave`
    /// * `support more than 1472 slots`
    ///
    pub fn new(slots: usize) -> LabelStorageArea {
        if slots > 1472 {
            todo!();
        }

        LabelStorageArea {
            label_index_blocks: [
                LabelIndexBlock::new(0, slots),
                LabelIndexBlock::new(1, slots),
            ],
            label_slots: Vec::with_capacity(slots),
            active_lib: -1, // Invalid
            current_seq: 0, // Invalid
        }
    }

    /// Converts an array of bytes into a Label Storage Area
    ///
    /// The Label Storage Area will be created from the bytes, even if there aren't valid label
    /// index blocks and/or labels. The only reason this function can "fail" is if there are the
    /// incorrect number of bytes provided.
    ///
    /// # Arguments
    ///
    /// * `bytes`: The bytes
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use cxl_rs::lsa::LabelStorageArea;
    /// use std::{io::Read, fs::File};
    ///
    /// let mut file = File::open("file")?;
    /// let mut buffer = Vec::new();
    /// file.read_to_end(&mut buffer)?;
    /// let lsa = LabelStorageArea::from_bytes(&buffer)?;
    /// # Ok::<(), std::io::Error>(())
    /// ```
    pub fn from_bytes(bytes: &[u8]) -> Result<Self, std::io::Error> {
        if bytes.len() < 1280 {
            return Err(Error::new(
                ErrorKind::UnexpectedEof,
                "LSA must be at least 1280 bytes",
            ));
        }
        if bytes.len() % 256 != 0 {
            return Err(Error::new(
                ErrorKind::InvalidInput,
                "LSA must be a multiple of 256 bytes",
            ));
        }

        let lib0: LabelIndexBlock = LabelIndexBlock::from_bytes(&bytes[..256]);
        let lib1: LabelIndexBlock = LabelIndexBlock::from_bytes(&bytes[256..512]);

        let mut labels: Vec<Label> = Vec::new();
        let mut index = 512;
        while index < bytes.len() {
            let label = label_from_bytes(&bytes[index..index + 256])?;
            labels.push(label);
            index += 256;
        }

        // CXL 2.0 spec states:
        //    When reading Label Index Blocks, software shall only consider index blocks valid when
        //    their Sig, MyOff, OtherOff, and Checksum fields are correct. In addition, any blocks
        //    with Seq set to zero are discarded as invalid. Finally, if more than 1 Label Index
        //    Block is found to be valid the one with the older sequence number (immediately
        //    counterclockwise to the other, according to Figure159 below) is discarded. If all
        //    checks pass and the sequence numbers match, the index block at the higher offset
        //    shall be considered the valid block.
        //
        // The follow sets active.0 to which label is active and active.1 to the sequence number
        // for that active label.
        //
        // This is the C code equivalent which is much more concise. It lets
        // static inline unsigned nd_inc_seq(unsigned seq)
        // {
        //    static const unsigned next[] = { 0, 2, 3, 1 };
        //    return next[seq & 3];
        // }
        //
        // static u32 best_seq(u32 a, u32 b)
        // {
        //     a &= NSINDEX_SEQ_MASK;
        //     b &= NSINDEX_SEQ_MASK;
        //
        //            if (a == 0 || a == b)
        //                return b;
        //            else if (b == 0)
        //                return a;
        //            else if (nd_inc_seq(a) == b)
        //                return b;
        //            else
        //                return a;
        //        }
        let active: (usize, u32) = match (lib0.valid().is_err(), lib1.valid().is_err()) {
            // Both blocks are invalid. The whole LSA is invalid.
            (true, true) => {
                println!("Couldn't find a valid Label Index Block");
                println!("\tInvalid LIB0: {:?}", lib0.valid());
                println!("\tInvalid LIB1: {:?}", lib1.valid());
                (2, 0)
            }
            // Only single valid label index block
            (false, true) => (0_usize, lib0.get_sequence()),
            (true, false) => (1_usize, lib1.get_sequence()),
            // If both are valid
            (false, false) => {
                match (lib0.get_sequence(), lib1.get_sequence()) {
                    // Second block wins when same sequence
                    (1, 1) | (2, 2) | (3, 3) => (1_usize, lib1.get_sequence()),
                    // Block 0 has a "higher" sequence
                    (2, 1) | (3, 2) | (1, 3) => (0_usize, lib0.get_sequence()),
                    // Block 1 has a "higher" sequence
                    (1, 2) | (2, 3) | (3, 1) => (1_usize, lib1.get_sequence()),
                    (a, b) => panic!("Unexpected sequence ({}, {})", a, b),
                }
            }
        };

        Ok(LabelStorageArea {
            label_index_blocks: [lib0, lib1],
            label_slots: labels,
            active_lib: active.0 as isize,
            current_seq: active.1,
        })
    }

    // This should probably be API at some point.
    //
    //    fn get_active_label_index_block(&self) -> Option<&LabelIndexBlock> {
    //        match self.active_lib {
    //            0 => Some(&self.label_index_blocks[0]),
    //            1 => Some(&self.label_index_blocks[1]),
    //            _ => None,
    //        }
    //    }
    //
    //    fn get_inactive_label_index_block(&self) -> Option<&LabelIndexBlock> {
    //        match self.active_lib {
    //            0 => Some(&self.label_index_blocks[1]),
    //            1 => Some(&self.label_index_blocks[0]),
    //            _ => None,
    //        }
    //    }
    //
    //    fn get_inactive_label_index_block_as_mut(&mut self) -> Option<&mut LabelIndexBlock> {
    //        match self.active_lib {
    //            0 => Some(&mut self.label_index_blocks[1]),
    //            1 => Some(&mut self.label_index_blocks[0]),
    //            _ => None,
    //        }
    //    }
    //

    pub fn push_label(&mut self, label: Label) -> Result<usize, std::io::Error> {
        self.label_slots.push(label);
        if self.label_slots.capacity() > 1472 {
            Err(Error::new(ErrorKind::Other, "No slots left"))
        } else {
            Ok(self.label_slots.len() - 1)
        }
    }

    /// Create a new region label in the next available slot. Returns the slot number.
    ///
    /// This is similar to [push_namespace_label](#method.push_namespace_label) and
    /// [push_empty_label](#method.push_empty_label).
    ///
    /// # Arguments
    ///
    /// * `size`: The size of the region to be created.
    ///
    /// # Examples
    ///
    /// ```
    /// use cxl_rs::lsa::LabelStorageArea;
    /// let mut lsa = LabelStorageArea::new(20);
    ///
    /// // Create a label of 4G, it should be slot 0
    /// let r1 = lsa.push_region_label(4 << 30)?;
    ///
    /// assert_eq!(r1, 0);
    /// # Ok::<(), std::io::Error>(())
    /// ```
    pub fn push_region_label(&mut self, size: usize) -> Result<usize, std::io::Error> {
        let mut region = RegionLabel::new().size(size);
        region.update_checksum();
        self.push_label(Region(region))
    }

    /// Create a new namespace label in the next available slot. Returns the slot number.
    ///
    /// This is similar to [push_region_label](#method.push_region_label) and
    /// [push_empty_label](#method.push_empty_label).
    ///
    /// # Arguments
    ///
    /// * `size`: The size of the namespace to be created.
    ///
    /// # Examples
    ///
    /// ```
    /// use cxl_rs::lsa::LabelStorageArea;
    /// let mut lsa = LabelStorageArea::new(20);
    ///
    /// // Create a label of 4G, it should be slot 0
    /// let n1 = lsa.push_namespace_label(4 << 30)?;
    ///
    /// assert_eq!(n1, 0);
    /// # Ok::<(), std::io::Error>(())
    /// ```
    pub fn push_namespace_label(&mut self, size: usize) -> Result<usize, std::io::Error> {
        let mut namespace = NamespaceLabel::new().size(size);
        namespace.update_checksum();
        self.push_label(Namespace(namespace))
    }

    /// Create an empty label and returns the slot number.
    ///
    /// This is similar to [push_region_label](#method.push_region_label) and
    /// [push_namespace_label](#method.push_namespace_label). The function isn't very useful and
    /// will probably be deprecated in the future.
    pub fn push_empty_label(&mut self) -> Result<usize, std::io::Error> {
        self.push_label(Empty)
    }

    /// Replaces a given slot with a label.
    pub fn replace_label_slot(&mut self, slot: usize, label: Label) -> Result<(), std::io::Error> {
        self.push_label(label)?;
        self.label_slots.swap_remove(slot);
        Ok(())
    }

    /// Marks label slots in use for the given label index block.
    ///
    /// In general it is unwise to do this for the current, active label index block.
    ///
    /// # Arguments
    ///
    /// * `lib`: The label index block to update.
    /// * `labels`: The labels to be marked as in use.
    ///
    /// # Examples
    ///
    /// ```
    /// use cxl_rs::lsa::LabelStorageArea;
    /// let mut lsa = LabelStorageArea::new(20);
    ///
    /// // Create a label of 4G, it should be slot 0
    /// let r1 = lsa.push_region_label(4 << 30)?;
    ///
    /// lsa.retain_label_slots(0, vec![r1]);
    /// # Ok::<(), std::io::Error>(())
    /// ```
    pub fn retain_label_slots(&mut self, lib: usize, labels: Vec<usize>) {
        for label in labels {
            self.label_index_blocks[lib].retain_slot(label);
        }
    }

    /// Sets the active label_index_block
    pub fn set_active(&mut self, lib: usize) {
        let new_seq = next_seq(self.current_seq);

        let label = match self.active_lib {
            0 => &mut self.label_index_blocks[1],
            1 => &mut self.label_index_blocks[0],
            _ => &mut self.label_index_blocks[0],
        };

        label.set_sequence(new_seq);
        label.update_checksum();

        self.current_seq = new_seq;
        self.active_lib = lib as isize;
    }

    fn write_label_index_block(
        &self,
        v: &mut Vec<u8>,
        which: usize,
    ) -> Result<usize, std::io::Error> {
        let lib = self.label_index_blocks[which];
        let view = &lib as *const _ as *const u8;
        let slice = unsafe { std::slice::from_raw_parts(view, 256) };
        v.write(slice)
    }

    fn write_labels(&self, v: &mut Vec<u8>) -> Result<usize, std::io::Error> {
        let mut n = 0;
        for label in &self.label_slots {
            // FIXME: can we make this not so ugly?
            let l = match label {
                Region(r) => r.as_bytes(),
                Namespace(n) => n.as_bytes(),
                Vendor(_v) => &[0; 256], //FIXME
                _ => &[0; 256],
            };
            n += match v.write(l) {
                Err(e) => return Err(e),
                Ok(_) => 256,
            }
        }
        Ok(n)
    }

    /// Return a byte vec of the LSA
    pub fn to_vec(&self) -> Result<Vec<u8>, std::io::Error> {
        let lib0 = &self.label_index_blocks[0];
        let lib1 = &self.label_index_blocks[1];
        let mut len = lib0.get_size();
        len += lib1.get_size();
        len += self.label_slots.len() * 256;

        let mut v: Vec<u8> = Vec::with_capacity(len);
        self.write_label_index_block(&mut v, 0)?;
        self.write_label_index_block(&mut v, 1)?;
        self.write_labels(&mut v)?;

        Ok(v)
    }

    /// Verify contents of the LSA.
    pub fn verify(&self) -> Result<(), std::io::Error> {
        if self.active_lib > 1 || self.active_lib < 0 {
            return Err(Error::new(ErrorKind::Other, "No active label found"));
        }

        let active = self.active_lib as usize;
        let index = &self.label_index_blocks[active];

        if let Err(e) = index.valid() {
            return Err(Error::new(ErrorKind::Other, e));
        }

        // TODO: Call verify on all active slots

        Ok(())
    }
}

impl Serialize for LabelStorageArea {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        let mut state =
            serializer.serialize_struct("LabelStorageArea", 1 + 2 + self.label_slots.len())?;
        state.serialize_field("Valid block", &self.active_lib)?;
        state.serialize_field("LBI0", &self.label_index_blocks[0])?;
        state.serialize_field("LBI1", &self.label_index_blocks[1])?;
        for (_i, label) in self.label_slots.iter().enumerate() {
            match label {
                Label::Region(region) => state.serialize_field("region", region),
                Label::Namespace(namespace) => state.serialize_field("namespace", namespace),
                Label::Vendor(raw) => state.serialize_field("vendor", raw),
                Label::Empty => state.serialize_field("empty", &0),
            }?;
        }
        state.end()
    }
}