use crc32fast::Hasher;
use lzma_rs::decompress::raw::Lzma2Decoder;
use std::{
    collections::VecDeque,
    io::{BufRead, BufReader, Error, Read, Result, Write},
};

#[derive(Debug)]
pub struct XzDecoder<R> {
    compressed_reader: BufReader<R>,
    stream_size: usize,
    buf: VecDeque<u8>,
    check_size: usize,
    records: Vec<(usize, usize)>,
    flags: [u8; 2],
}

impl<R: Read> XzDecoder<R> {
    pub fn new(inner: R) -> Self {
        XzDecoder {
            compressed_reader: BufReader::new(inner),
            stream_size: 0,
            buf: VecDeque::new(),
            check_size: 0,
            records: vec![],
            flags: [0, 0],
        }
    }
}

struct CountReader<'a, R: BufRead> {
    inner: &'a mut R,
    count: &'a mut usize,
}

impl<R: BufRead> Read for CountReader<'_, R> {
    fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
        let count = self.inner.read(buf)?;
        *self.count += count;
        Ok(count)
    }
}

impl<R: BufRead> BufRead for CountReader<'_, R> {
    fn fill_buf(&mut self) -> Result<&[u8]> {
        self.inner.fill_buf()
    }

    fn consume(&mut self, amt: usize) {
        self.inner.consume(amt);
        *self.count += amt;
    }
}

struct BufWriter<'a> {
    inner: &'a mut [u8],
    written: &'a mut usize,
    total: &'a mut usize,
    rest: &'a mut VecDeque<u8>,
}

impl<'a> Write for BufWriter<'a> {
    fn write(&mut self, buf: &[u8]) -> Result<usize> {
        if self.inner.len() > *self.written {
            let len = std::cmp::min(buf.len(), self.inner.len() - *self.written);
            self.inner[*self.written..*self.written + len].copy_from_slice(&buf[..len]);
            *self.written += len;
            *self.total += len;
            Ok(len)
        } else {
            self.rest.extend(buf.iter());
            *self.total += buf.len();
            Ok(buf.len())
        }
    }

    fn flush(&mut self) -> Result<()> {
        Ok(())
    }
}

fn error<T>(s: &'static str) -> Result<T> {
    Err(Error::new(std::io::ErrorKind::InvalidData, s))
}

fn get_multibyte<R: Read>(input: &mut R, hasher: &mut Hasher) -> Result<u64> {
    let mut result = 0;
    for i in 0..9 {
        let mut b = [0u8; 1];
        input.read_exact(&mut b)?;
        hasher.update(&b);
        let b = b[0];
        result ^= ((b & 0x7F) as u64) << (i * 7);
        if (b & 0x80) == 0 {
            return Ok(result);
        }
    }
    error("Invalid multi-byte encoding")
}

impl<R: Read> Read for XzDecoder<R> {
    fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
        if !self.buf.is_empty() {
            let len = std::cmp::min(buf.len(), self.buf.len());
            buf[..len].copy_from_slice(&self.buf.as_slices().0[..len]);
            self.buf.drain(..len);
            return Ok(len);
        }
        let mut reader = CountReader {
            inner: &mut self.compressed_reader,
            count: &mut self.stream_size,
        };
        if *reader.count == 0 {
            let mut b = [0u8; 12];
            match reader.read(&mut b) {
                Ok(0) => return Ok(0),
                Err(e) => return Err(e),
                _ => (),
            }
            if b[..6] != b"\xFD7zXZ\0"[..] {
                return error("Invalid XZ header");
            }
            self.flags = [b[6], b[7]];
            if self.flags[0] != 0 || self.flags[1] & 0xF0 != 0 {
                return error("Invalid XZ stream flags");
            }
            match self.flags[1] & 0x0F {
                0 => self.check_size = 0,
                1 => self.check_size = 4,
                _ => return error("Unsupported XZ stream flags"),
            }
            let mut digest = Hasher::new();
            digest.update(&self.flags);
            if digest.finalize().to_le_bytes() != b[8..] {
                return error("Invalid XZ stream flags CRC32");
            }
        }

        let block_begin = *reader.count;
        let mut b = [0u8; 1];
        reader.read_exact(&mut b)?;

        let mut digest = Hasher::new();
        digest.update(&b);
        if b[0] == 0 {
            // index
            let num_records = get_multibyte(&mut reader, &mut digest)?;
            if num_records != self.records.len() as u64 {
                return error("Invalid XZ index record count");
            }
            for (unpadded_size, total) in &self.records {
                if get_multibyte(&mut reader, &mut digest)? != *unpadded_size as u64 {
                    return error("Invalid XZ unpadded size");
                }
                if get_multibyte(&mut reader, &mut digest)? != *total as u64 {
                    return error("Invalid XZ uncompressed size");
                }
            }
            let mut size = *reader.count - block_begin;
            let mut b = vec![0u8; (4 - (size & 0x3)) & 0x3];
            reader.read_exact(b.as_mut_slice())?;
            if !b.iter().all(|&b| b == 0) {
                return error("Invalid XZ index padding");
            }
            digest.update(b.as_slice());
            size += b.len();
            let mut b = [0u8; 16];
            reader.read_exact(&mut b)?;
            if digest.finalize().to_le_bytes() != b[..4] {
                return error("Invalid XZ index CRC32");
            }
            let mut digest = Hasher::new();
            digest.update(&b[8..14]);
            if digest.finalize().to_le_bytes() != b[4..8] {
                return error("Invalid XZ footer CRC32");
            }
            if b[8..12] != ((size >> 2) as u32).to_le_bytes() {
                return error("Invalid XZ footer size");
            }
            if self.flags != b[12..14] {
                return error("Invalid XZ footer flags");
            }
            if &b[14..16] != b"YZ" {
                return error("Invalid XZ footer magic");
            }
            let mut b = vec![0u8; (4 - (*reader.count & 0x3)) & 0x3];
            reader.read_exact(b.as_mut_slice())?;
            if !b.iter().all(|&b| b == 0) {
                return error("Invalid XZ footer padding");
            }
            *reader.count = 0;
            return self.read(buf);
        }

        // block
        let header_end = ((b[0] as usize) << 2) - 1 + *reader.count;
        let mut b = [0u8; 1];
        reader.read_exact(&mut b)?;
        digest.update(&b);
        let flags = b[0];
        let num_filters = (flags & 0x03) + 1;

        if flags & 0x3C != 0 {
            return error("Invalid XZ block flags");
        }
        if flags & 0x40 != 0 {
            get_multibyte(&mut reader, &mut digest)?;
        }
        if flags & 0x80 != 0 {
            get_multibyte(&mut reader, &mut digest)?;
        }
        for _ in 0..num_filters {
            let filter_id = get_multibyte(&mut reader, &mut digest)?;
            if filter_id != 0x21 {
                return error("Unsupported XZ filter ID");
            }
            let properties_size = get_multibyte(&mut reader, &mut digest)?;
            if properties_size != 1 {
                return error("Unsupported XZ filter properties size");
            }
            reader.read_exact(&mut b)?;
            if b[0] & 0xC0 != 0 {
                return error("Unsupported XZ filter properties");
            }
            digest.update(&b);
        }
        let mut b = vec![0u8; header_end - *reader.count];
        reader.read_exact(b.as_mut_slice())?;
        if !b.iter().all(|&b| b == 0) {
            return error("Invalid XZ block header padding");
        }
        digest.update(b.as_slice());

        let mut b = [0u8; 4];
        reader.read_exact(&mut b)?;
        if digest.finalize().to_le_bytes() != b {
            return error("Invalid XZ block header CRC32");
        }
        let mut written = 0;
        let mut total = 0;
        Lzma2Decoder::new().decompress(
            &mut reader,
            &mut BufWriter {
                inner: buf,
                written: &mut written,
                rest: &mut self.buf,
                total: &mut total,
            },
        )?;

        let unpadded_size = *reader.count - block_begin;
        self.records.push((unpadded_size, total));
        // ignore check here since zip itself will check it
        let mut b = vec![0u8; ((4 - (unpadded_size & 0x3)) & 0x3) + self.check_size];
        reader.read_exact(b.as_mut_slice())?;
        if !b.as_slice()[..self.check_size].iter().all(|&b| b == 0) {
            return error("Invalid XZ block padding");
        }
        Ok(written)
    }
}

impl<R: Read> XzDecoder<R> {
    pub fn into_inner(self) -> R {
        self.compressed_reader.into_inner()
    }
}