ttb 0.1.1

use {
    regenboog::RgbaU8,
    std::io::{self, Write},
};

use crate::{MAGIC_NUMBER, VERSION};

#[derive(Debug, PartialEq)]
enum State {
    Header,
    Blocks,
}

/// Wraps a type implementing `Write` and can be used to write in the TTB format.
#[must_use = "the `finish` method must be called to correctly end the stream and write all data"]
pub struct Writer<W: Write> {
    state: State,
    palette_size: u8,
    brick_types: Vec<String>,
    pending_brick_types: Vec<u16>,
    pending_bricks: Vec<(u16, i32, i32, i32, u8, u8)>,
    inner: W,
}

impl<W: Write> Writer<W> {
    /// Create a new `Writer` from a type implementing `Write`.
    pub fn new(writer: W) -> Self {
        Self {
            state: State::Header,
            palette_size: 0,
            brick_types: Vec::new(),
            pending_brick_types: Vec::new(),
            pending_bricks: Vec::new(),
            inner: writer,
        }
    }

    /// Extract the underlying type implementing `Write`, destroying this `Writer`.
    pub fn into_inner(self) -> W {
        self.inner
    }

    /// Write the TTB header, this contains the build's name and description and the palette.
    ///
    /// # Panics
    ///
    /// If this `Writer` has already written this header, this will panic.
    ///
    /// The byte length of `build_name` must be less than 256, else this will panic.
    ///
    /// The byte length of `build_description` must be less than 65536, else this will panic.
    pub fn write_header(
        &mut self,
        build_name: &str,
        build_description: &str,
        palette: &[RgbaU8],
    ) -> io::Result<()> {
        if self.state != State::Header {
            panic!(
                "invalid state: expecting state Header but am in state {:?}",
                self.state
            );
        }
        assert!(build_name.len() < 256);
        assert!(build_description.len() < 65536);
        self.inner.write(MAGIC_NUMBER)?;
        self.inner.write(&[VERSION])?;
        let build_name_len = build_name.len() as u8;
        let build_description_len = build_description.len() as u16;
        self.inner.write(&[build_name_len])?;
        self.inner.write(build_name.as_bytes())?;
        self.inner.write(&build_description_len.to_be_bytes())?;
        self.inner.write(build_description.as_bytes())?;
        let palette_size = palette.len() as u8;
        self.inner.write(&[palette_size])?;
        for &col in palette {
            let col_arr: [u8; 4] = col.into();
            self.inner.write(&col_arr)?;
        }
        self.palette_size = palette_size;
        self.state = State::Blocks;
        Ok(())
    }

    /// Write a brick's data, which possibly will also add its type to the brick type table.
    ///
    /// # Panics
    ///
    /// If the header has not yet been written, this will panic.
    ///
    /// If the color is greater than or equal to the palette's size, this will panic.
    ///
    /// If the orientation is greater than or equal to 4, this will panic.
    ///
    /// If the brick's type name is longer than 255 bytes, this will panic.
    pub fn write_brick(
        &mut self,
        brick_type_name: &str,
        x: i32,
        y: i32,
        z: i32,
        orientation: u8,
        color: u8,
    ) -> io::Result<()> {
        if self.state != State::Blocks {
            panic!(
                "invalid state: expecting state Blocks but am in state {:?}",
                self.state
            );
        }
        if color >= self.palette_size {
            panic!("color overflows palette: {}", color);
        }
        if orientation >= 4 {
            panic!("invalid brick orientation: {}", orientation);
        }
        assert!(brick_type_name.len() < 256);
        let tyid = if let Some(tyid) = self.brick_types.iter().position(|x| x == brick_type_name) {
            tyid as u16
        } else {
            let tyid = self.brick_types.len() as u16;
            self.brick_types.push(brick_type_name.to_owned());
            self.add_brick_type(tyid)?;
            tyid
        };
        self.add_brick(tyid, x, y, z, orientation, color)?;
        Ok(())
    }

    /// Finish writing, this must be called after all write operations have been done.
    pub fn finish(mut self) -> io::Result<()> {
        self.flush_brick_types()?;
        self.flush_bricks()?;
        Ok(())
    }

    fn add_brick_type(&mut self, tyid: u16) -> io::Result<()> {
        if self.pending_brick_types.len() >= 255 {
            self.flush_brick_types()?;
        }
        self.pending_brick_types.push(tyid);
        Ok(())
    }

    fn add_brick(
        &mut self,
        tyid: u16,
        x: i32,
        y: i32,
        z: i32,
        orientation: u8,
        color: u8,
    ) -> io::Result<()> {
        if self.pending_bricks.len() >= 255 {
            self.flush_brick_types()?;
            self.flush_bricks()?;
        }
        self.pending_bricks
            .push((tyid, x, y, z, orientation, color));
        Ok(())
    }

    fn flush_brick_types(&mut self) -> io::Result<()> {
        if !self.pending_brick_types.is_empty() {
            let len = self.pending_brick_types.len();
            self.inner.write_all(&[1, len as u8])?;
            for idx in self.pending_brick_types.drain(..) {
                let brick_type_name = &self.brick_types[idx as usize];
                self.inner.write_all(&[brick_type_name.len() as u8])?;
                self.inner.write_all(brick_type_name.as_bytes())?;
            }
        }
        Ok(())
    }

    fn flush_bricks(&mut self) -> io::Result<()> {
        if !self.pending_bricks.is_empty() {
            let len = self.pending_bricks.len();
            self.inner.write_all(&[2, len as u8])?;
            for (tyid, x, y, z, orientation, color) in self.pending_bricks.drain(..) {
                self.inner.write_all(&tyid.to_be_bytes())?;
                self.inner.write_all(&x.to_be_bytes())?;
                self.inner.write_all(&y.to_be_bytes())?;
                self.inner.write_all(&z.to_be_bytes())?;
                self.inner.write_all(&[orientation, color])?;
            }
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use {regenboog::RgbaU8, std::io::Cursor};

    use super::*;

    #[test]
    #[should_panic = "invalid state: expecting state Header but am in state Blocks"]
    fn ttb_writer_header_doesnt_write_twice() {
        let mut buf = [0; 2048];
        let cursor = Cursor::new(&mut buf[..]);
        let mut writer = Writer::new(cursor);
        writer
            .write_header("test", "test", &[RgbaU8::RED, RgbaU8::GREEN, RgbaU8::BLUE])
            .unwrap();
        writer
            .write_header("test", "test", &[RgbaU8::RED, RgbaU8::GREEN, RgbaU8::BLUE])
            .unwrap();
    }

    #[test]
    #[should_panic = "invalid state: expecting state Blocks but am in state Header"]
    fn ttb_writer_header_must_be_written() {
        let mut buf = [0; 2048];
        let cursor = Cursor::new(&mut buf[..]);
        let mut writer = Writer::new(cursor);
        writer.write_brick("1x1 Plate", 0, 0, 0, 0, 0).unwrap();
    }

    #[test]
    #[should_panic = "color overflows palette: 10"]
    fn ttb_writer_color_overflows_palette() {
        let mut buf = [0; 2048];
        let cursor = Cursor::new(&mut buf[..]);
        let mut writer = Writer::new(cursor);
        writer
            .write_header("test", "test", &[RgbaU8::RED, RgbaU8::GREEN, RgbaU8::BLUE])
            .unwrap();
        writer.write_brick("1x1 Plate", 0, 0, 0, 0, 10).unwrap();
    }

    #[test]
    #[should_panic = "invalid brick orientation: 5"]
    fn ttb_writer_invalid_orienbtation() {
        let mut buf = [0; 2048];
        let cursor = Cursor::new(&mut buf[..]);
        let mut writer = Writer::new(cursor);
        writer
            .write_header("test", "test", &[RgbaU8::RED, RgbaU8::GREEN, RgbaU8::BLUE])
            .unwrap();
        writer.write_brick("1x1 Plate", 0, 0, 0, 5, 2).unwrap();
    }
}