rustyjson

use flate2::write::GzEncoder;
use flate2::Compression;
use rustler::NifUnitEnum;
use std::io::{BufWriter, Error, Write};

#[derive(NifUnitEnum)]
pub enum Algs {
    None,
    Gzip,
}

/// Writer that can be converted to the final output bytes
pub enum Writer {
    Plain(Vec<u8>),
    Gzip(BufWriter<GzEncoder<Vec<u8>>>),
}

impl Write for Writer {
    fn write(&mut self, buf: &[u8]) -> Result<usize, Error> {
        match self {
            Writer::Plain(v) => v.write(buf),
            Writer::Gzip(w) => w.write(buf),
        }
    }

    fn flush(&mut self) -> Result<(), Error> {
        match self {
            Writer::Plain(v) => v.flush(),
            Writer::Gzip(w) => w.flush(),
        }
    }

    fn write_all(&mut self, buf: &[u8]) -> Result<(), Error> {
        match self {
            Writer::Plain(v) => v.write_all(buf),
            Writer::Gzip(w) => w.write_all(buf),
        }
    }
}

impl Writer {
    /// Get the final output buffer, consuming self
    pub fn get_buf(self) -> Result<Vec<u8>, Error> {
        match self {
            Writer::Plain(v) => Ok(v),
            Writer::Gzip(mut w) => {
                w.flush()?;
                let encoder = w.into_inner().map_err(|e| e.into_error())?;
                let vec = encoder.finish()?;
                Ok(vec)
            }
        }
    }
}

pub fn get_writer(opts: Option<(Algs, Option<u32>)>) -> Writer {
    match opts {
        Some((Algs::Gzip, None)) => Writer::Gzip(BufWriter::with_capacity(
            10_240,
            GzEncoder::new(Vec::with_capacity(4096), Compression::default()),
        )),
        Some((Algs::Gzip, Some(lv))) => Writer::Gzip(BufWriter::with_capacity(
            10_240,
            GzEncoder::new(Vec::with_capacity(4096), Compression::new(lv)),
        )),
        _ => Writer::Plain(Vec::with_capacity(4096)),
    }
}