zeekstd_cli 0.4.5

use std::{fs::File, io::Write};

use anyhow::{Context, Result, anyhow};
use indicatif::{ProgressBar, ProgressDrawTarget, ProgressStyle};
use zeekstd::{DecodeOptions, Decoder, SeekTable};
use zstd_safe::{DCtx, DParameter};

use crate::args::DecompressArgs;

pub struct Decompressor<'a> {
    decoder: Decoder<'a, File>,
    bar: Option<ProgressBar>,
}

impl Decompressor<'_> {
    pub fn new(
        args: &DecompressArgs,
        prefix_len: Option<u64>,
        progress_style: Option<ProgressStyle>,
    ) -> Result<Self> {
        let mut src = File::open(&args.input_file).context("Failed to open input file")?;
        let seek_table = match &args.common.seek_table_file {
            Some(path) => {
                let mut file = File::open(path).context("Failed to open seek table file")?;
                SeekTable::from_reader(&mut file)
            }
            None => SeekTable::from_seekable(&mut src),
        }
        .context("Failed to parse seek table")?;

        let offset = args
            .offset(&seek_table)
            .context("Failed to get decompression offset")?;

        let offset_limit = args
            .offset_limit(&seek_table)
            .context("Failed to get decompression offset limit")?;

        let bar = if let Some(style) = progress_style {
            let bar = ProgressBar::with_draw_target(
                Some(offset_limit),
                ProgressDrawTarget::stderr_with_hz(5),
            )
            .with_style(style);

            bar.set_position(offset);

            Some(bar)
        } else {
            None
        };

        let mut dctx = DCtx::try_create().context("Failed to create decompression context")?;
        if let Some(len) = prefix_len {
            let window_log_max = if len == 0 { 0 } else { len.ilog2() + 1 };
            dctx.set_parameter(DParameter::WindowLogMax(window_log_max))
                .map_err(|c| {
                    anyhow!(
                        "Failed to set max window log: {}",
                        zstd_safe::get_error_name(c)
                    )
                })?;
        }

        let decoder = DecodeOptions::with_dctx(src, dctx)
            .seek_table(seek_table)
            .offset(offset)
            .offset_limit(offset_limit)
            .into_decoder()
            .context("Failed to create decoder")?;

        Ok(Self { decoder, bar })
    }
}

impl<'a> Decompressor<'a> {
    pub fn decompress_into<'b: 'a, W: Write>(
        mut self,
        writer: &mut W,
        prefix: Option<&'b [u8]>,
    ) -> Result<u64> {
        let mut buf = vec![0; DCtx::out_size()];
        let mut buf_pos = 0;
        let mut written = 0;

        loop {
            let n = self
                .decoder
                .decompress_with_prefix(&mut buf[buf_pos..], prefix)
                .context("Failed to decompress data")?;
            if n == 0 {
                break;
            }
            if let Some(bar) = &self.bar {
                bar.inc(n as u64);
            }
            buf_pos += n;
            if buf_pos == buf.len() {
                writer
                    .write_all(&buf)
                    .context("Failed to write decompressed data")?;
                written += buf_pos as u64;
                buf_pos = 0;
            }
        }

        writer
            .write_all(&buf[..buf_pos])
            .context("Failed to write decompressed data")?;
        written += buf_pos as u64;

        if let Some(bar) = &self.bar {
            bar.finish_and_clear();
        }

        Ok(written)
    }
}