use core::fmt;

#[derive(Clone, Copy, Debug, Default)]
pub struct FormatSizeOptions {}

pub struct FormatSize(u64, FormatSizeOptions);

pub struct CharsWriter<'a> {
    buffer: &'a mut [u8],
    pos: usize,
}

impl<'a> CharsWriter<'a> {
    pub const fn new(buffer: &'a mut [u8]) -> Self {
        Self { buffer, pos: 0 }
    }

    pub fn push_char(&mut self, ch: char) -> bool {
        let w = ch.len_utf8();
        if self.pos + w < self.buffer.len() {
            ch.encode_utf8(&mut self.buffer[self.pos..]);
            self.pos += w;
            true
        } else {
            false
        }
    }

    pub fn as_str(&self) -> &str {
        // Safe, the implementation only allows full valid chars to be inserted
        unsafe { core::str::from_utf8_unchecked(&self.buffer[..self.pos]) }
    }
}

impl fmt::Write for CharsWriter<'_> {
    fn write_str(&mut self, s: &str) -> fmt::Result {
        for ch in s.chars() {
            if !self.push_char(ch) {
                break;
            }
        }
        Ok(())
    }
}

impl FormatSize {
    pub fn default(size: u64) -> Self {
        Self(size, FormatSizeOptions::default())
    }
}

impl fmt::Display for FormatSize {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        use fmt::Write;

        let mut buffer = [0; 128];
        let mut writer = CharsWriter::new(&mut buffer);
        let ilog2 = self.0.checked_ilog2().unwrap_or(0);
        let (dot, div, suffix) = match ilog2 / 10 {
            0 => (false, 1, "B"),
            1 => (true, 1024, " KiB"),
            2 => (true, 1024 * 1024, " MiB"),
            3 => (true, 1024 * 1024 * 1024, " GiB"),
            _ => (true, 1024 * 1024 * 1024 * 1024, " TiB"),
        };
        let integer = self.0 / div;
        if dot && integer < 100 {
            let fractional = (((self.0 * 1000) / div) % 1000) / 10;
            write!(writer, "{integer}.{fractional:02}{suffix}")?;
        } else {
            write!(writer, "{integer}{suffix}")?;
        }

        let chars = writer.as_str();
        fmt::Display::fmt(chars, f)
    }
}

#[cfg(test)]
mod tests {
    use alloc::format;

    use crate::fmt::FormatSize;

    #[test]
    fn test_fmt_size() {
        assert_eq!(format!("{}", FormatSize::default(1023)), "1023B");
        assert_eq!(format!("{}", FormatSize::default(1024)), "1.00 KiB");
        assert_eq!(format!("{}", FormatSize::default(1025)), "1.00 KiB");
        assert_eq!(format!("{}", FormatSize::default(1111)), "1.08 KiB");
        assert_eq!(format!("{}", FormatSize::default(5000)), "4.88 KiB");
        assert_eq!(format!("{}", FormatSize::default(5000000)), "4.76 MiB");
        assert_eq!(format!("{}", FormatSize::default(5000000000)), "4.65 GiB");
        assert_eq!(
            format!("{}", FormatSize::default(5000000000000)),
            "4.54 TiB"
        );
    }
}
