thyme/src/file.rs

use crate::data::Mapping;
use crate::errors::{ThymeFileError, ThymeFileStructureError};
use crate::image::{Image, Pixel};

use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::io::{Read, Seek, Write};

const CURRENT_VERSION: &'static str = "0";

#[derive(Debug, Clone, Copy)]
pub enum StitchType {
    Normal,
    HalfUp,
    HalfDown,
}

impl StitchType {
    fn to_u8(&self) -> u8 {
        match self {
            StitchType::Normal => 'x' as u8,
            StitchType::HalfUp => '/' as u8,
            StitchType::HalfDown => '\\' as u8,
        }
    }
}

#[derive(Debug, Clone)]
pub struct Stitch {
    pub color: ColorIdx,
    pub typ: StitchType,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct Color {
    pub name: String,
    pub symbol: String,
    pub color: Pixel,
}

#[derive(Debug, Clone, Copy)]
pub struct ColorIdx {
    pub idx: u32,
}

#[derive(Debug)]
pub struct ThymeFile {
    pub width: u32,
    pub height: u32,
    pub palette: Vec<Color>,
    pub payload: Vec<Option<Stitch>>,
    pub metadata: HashMap<String, String>,
}

impl ThymeFile {
    pub fn from_image_and_config(image: &Image, Mapping(mapping): &Mapping) -> ThymeFile {
        let width = image.width;
        let height = image.height;

        let mut palette = Vec::new();
        let mut lookup = HashMap::new();
        let mut next_idx = 0;
        for (color, info) in mapping.iter() {
            if let Some(i) = info {
                let idx = ColorIdx { idx: next_idx };
                next_idx += 1;
                palette.push(Color {
                    color: *color,
                    name: i.name.to_owned(),
                    symbol: i.symbol.to_owned(),
                });
                lookup.insert(color, idx);
            }
        }

        let payload = image
            .iter()
            .map(|(_idx, pixel)| {
                if let Some(color) = lookup.get(&pixel) {
                    Some(Stitch {
                        color: *color,
                        typ: StitchType::Normal,
                    })
                } else {
                    None
                }
            })
            .collect();

        let metadata = HashMap::new();

        ThymeFile {
            width,
            height,
            palette,
            payload,
            metadata,
        }
    }

    pub fn blank(width: u32, height: u32) -> ThymeFile {
        let palette = Vec::new();
        let mut payload = Vec::with_capacity((width * height) as usize);
        for _ in 0..(width * height) {
            payload.push(None);
        }

        let metadata = HashMap::new();

        ThymeFile {
            width,
            height,
            palette,
            payload,
            metadata,
        }
    }

    pub fn from_stream<R>(stream: &mut R) -> Result<ThymeFile, ThymeFileError>
    where
        R: Read + Seek,
    {
        fn missing_component(component: &str) -> ThymeFileError {
            ThymeFileStructureError::new(format!(
                "Missing component in thyme file: `{}`",
                component
            ))
            .into()
        }
        let mut zip = zip::ZipArchive::new(stream)?;
        let mut width = None;
        let mut height = None;
        let mut metadata = None;
        let mut palette = None;
        let mut payload = None;
        for i in 0..zip.len() {
            let file = zip.by_index(i)?;
            match file.name() {
                "THYME_VERSION" => {
                    // TODO: actually check this
                }
                "dimensions.json" => {
                    let (w, h): (u32, u32) = serde_json::from_reader(file)?;
                    width = Some(w);
                    height = Some(h);
                }
                "metadata.json" => metadata = Some(serde_json::from_reader(file)?),
                "palette.json" => palette = Some(serde_json::from_reader(file)?),
                "payload.json" => {
                    let ps: Vec<Option<IntermediateStitch>> = serde_json::from_reader(file)?;
                    let ps: Result<Vec<Option<Stitch>>, ThymeFileError> = ps
                        .iter()
                        .map(|n| n.map(IntermediateStitch::to_stitch).transpose())
                        .collect();
                    payload = Some(ps?);
                }
                name => {
                    return Err(ThymeFileStructureError::new(format!(
                        "Unrecognized element: {}",
                        name
                    ))
                    .into())
                }
            }
            // whatever
        }
        Ok(ThymeFile {
            width: width.ok_or_else(|| missing_component("dimensions"))?,
            height: height.ok_or_else(|| missing_component("dimensions"))?,
            palette: palette.ok_or_else(|| missing_component("palette"))?,
            payload: payload.ok_or_else(|| missing_component("payload"))?,
            metadata: metadata.ok_or_else(|| missing_component("metadata"))?,
        })
    }

    pub fn to_stream<W>(&self, stream: &mut W) -> std::io::Result<()>
    where
        W: Write + Seek,
    {
        let mut zip = zip::ZipWriter::new(stream);
        let options =
            zip::write::FileOptions::default().compression_method(zip::CompressionMethod::Bzip2);

        // add version metadata
        zip.start_file("THYME_VERSION", options)?;
        writeln!(zip, "{}", CURRENT_VERSION)?;

        // add dimensions
        zip.start_file("dimensions.json", options)?;
        serde_json::to_writer(&mut zip, &(self.width, self.height))?;

        // add metadata
        zip.start_file("metadata.json", options)?;
        serde_json::to_writer(&mut zip, &self.metadata)?;

        // add palette
        zip.start_file("palette.json", options)?;
        serde_json::to_writer(&mut zip, &self.palette)?;

        // add image payload
        zip.start_file("payload.json", options)?;
        serde_json::to_writer(&mut zip, &self.json_stitches())?;

        // add version
        zip.finish()?;
        Ok(())
    }

    fn json_stitches(&self) -> Vec<Option<IntermediateStitch>> {
        self.payload
            .iter()
            .map(|stitch| match stitch {
                Some(s) => Some(IntermediateStitch(s.typ.to_u8(), s.color.idx)),
                None => None,
            })
            .collect()
    }

    pub fn iter<'a>(&'a self) -> ThymeImageIterator<'a> {
        ThymeImageIterator {
            palette: &self.palette,
            x: 0,
            y: 0,
            width: self.width,
            values: self.payload.iter(),
        }
    }
}

// internal structs for serializing/deserializing the image part
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
struct IntermediateStitch(u8, u32);

impl IntermediateStitch {
    fn to_stitch(self) -> Result<Stitch, ThymeFileError> {
        let typ = match self.0 {
            b'x' => StitchType::Normal,
            b'/' => StitchType::HalfUp,
            b'\\' => StitchType::HalfDown,
            _ => {
                return Err(ThymeFileStructureError::new(format!(
                    "Unknown stitch type: {}",
                    self.0
                ))
                .into())
            }
        };
        Ok(Stitch {
            typ,
            color: ColorIdx { idx: self.1 },
        })
    }
}

pub struct ThymeImageIterator<'a> {
    width: u32,
    x: u32,
    y: u32,
    palette: &'a [Color],
    values: std::slice::Iter<'a, Option<Stitch>>,
}

impl<'a> Iterator for ThymeImageIterator<'a> {
    type Item = ((u32, u32), Option<(StitchType, &'a Color)>);

    fn next(&mut self) -> Option<Self::Item> {
        let px = self.values.next();
        if let Some(px) = px {
            let r = px.as_ref().map(|stitch| {
                (stitch.typ, &self.palette[stitch.color.idx as usize])
            });
            let ret = ((self.x, self.y), r);
            self.x += 1;
            if self.x == self.width {
                self.x = 0;
                self.y += 1;
            }
            Some(ret)
        } else {
            None
        }
    }
}