use std::fmt;
use std::io::Read;

pub const MAX_DECOMPRESSED_SIZE: usize = 32 * 1024 * 1024;

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum CompressionError {
    OutputTooLarge,
    TruncatedInput,
    InvalidOffset,
    Decoder(String),
}

impl fmt::Display for CompressionError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::OutputTooLarge => formatter.write_str("decompressed payload exceeds the limit"),
            Self::TruncatedInput => formatter.write_str("compressed payload is truncated"),
            Self::InvalidOffset => formatter.write_str("compressed payload has an invalid offset"),
            Self::Decoder(message) => formatter.write_str(message),
        }
    }
}

impl std::error::Error for CompressionError {}

pub fn decompress_payload(source: &[u8]) -> Result<Vec<u8>, CompressionError> {
    if source.starts_with(&[0x28, 0xb5, 0x2f, 0xfd]) {
        let decoder = zstd::stream::read::Decoder::new(source)
            .map_err(|error| CompressionError::Decoder(error.to_string()))?;
        return read_bounded(decoder);
    }
    if source.starts_with(&[0x04, 0x22, 0x4d, 0x18]) {
        return read_bounded(lz4_flex::frame::FrameDecoder::new(source));
    }
    decompress_lz4_block(source, MAX_DECOMPRESSED_SIZE)
}

pub fn compress_lz4_block(source: &[u8]) -> Vec<u8> {
    lz4_flex::block::compress(source)
}

fn read_bounded(reader: impl Read) -> Result<Vec<u8>, CompressionError> {
    let mut output = Vec::new();
    reader
        .take((MAX_DECOMPRESSED_SIZE + 1) as u64)
        .read_to_end(&mut output)
        .map_err(|error| CompressionError::Decoder(error.to_string()))?;
    if output.len() > MAX_DECOMPRESSED_SIZE {
        return Err(CompressionError::OutputTooLarge);
    }
    Ok(output)
}

fn decompress_lz4_block(source: &[u8], maximum: usize) -> Result<Vec<u8>, CompressionError> {
    let mut output = Vec::new();
    let mut cursor = 0usize;

    while cursor < source.len() {
        let token = source[cursor];
        cursor += 1;

        let literal_length = read_length(source, &mut cursor, (token >> 4) as usize)?;
        if literal_length > source.len().saturating_sub(cursor) {
            return Err(CompressionError::TruncatedInput);
        }
        ensure_capacity(output.len(), literal_length, maximum)?;
        output.extend_from_slice(&source[cursor..cursor + literal_length]);
        cursor += literal_length;

        if cursor == source.len() {
            break;
        }
        if source.len().saturating_sub(cursor) < 2 {
            return Err(CompressionError::TruncatedInput);
        }

        let offset = usize::from(source[cursor]) | (usize::from(source[cursor + 1]) << 8);
        cursor += 2;
        if offset == 0 || offset > output.len() {
            return Err(CompressionError::InvalidOffset);
        }

        let match_length = read_length(source, &mut cursor, usize::from(token & 0x0f))?
            .checked_add(4)
            .ok_or(CompressionError::OutputTooLarge)?;
        ensure_capacity(output.len(), match_length, maximum)?;
        for _ in 0..match_length {
            output.push(output[output.len() - offset]);
        }
    }

    Ok(output)
}

fn read_length(
    source: &[u8],
    cursor: &mut usize,
    initial: usize,
) -> Result<usize, CompressionError> {
    if initial != 15 {
        return Ok(initial);
    }
    let mut length = initial;
    loop {
        let byte = *source
            .get(*cursor)
            .ok_or(CompressionError::TruncatedInput)?;
        *cursor += 1;
        length = length
            .checked_add(usize::from(byte))
            .ok_or(CompressionError::OutputTooLarge)?;
        if byte != 255 {
            return Ok(length);
        }
    }
}

fn ensure_capacity(
    current: usize,
    additional: usize,
    maximum: usize,
) -> Result<(), CompressionError> {
    if additional > maximum.saturating_sub(current) {
        return Err(CompressionError::OutputTooLarge);
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::{compress_lz4_block, decompress_lz4_block, CompressionError};

    #[test]
    fn lz4_block_round_trip() {
        let source = b"synthetic synthetic synthetic synthetic";
        let compressed = compress_lz4_block(source);
        assert_eq!(decompress_lz4_block(&compressed, 1024).unwrap(), source);
    }

    #[test]
    fn lz4_limit_is_enforced() {
        let compressed = compress_lz4_block(b"0123456789");
        assert_eq!(
            decompress_lz4_block(&compressed, 4).unwrap_err(),
            CompressionError::OutputTooLarge
        );
    }
}
