// spell-checker:ignore compressobj decompressobj zdict chunksize zlibmodule miniz

pub(crate) use zlib::make_module;

#[pymodule]
mod zlib {
    use crate::vm::{
        builtins::{PyBaseExceptionRef, PyBytes, PyBytesRef, PyIntRef, PyTypeRef},
        common::lock::PyMutex,
        convert::TryFromBorrowedObject,
        function::{ArgBytesLike, ArgPrimitiveIndex, ArgSize, OptionalArg},
        PyObject, PyPayload, PyResult, VirtualMachine,
    };
    use adler32::RollingAdler32 as Adler32;
    use crossbeam_utils::atomic::AtomicCell;
    use flate2::{
        write::ZlibEncoder, Compress, Compression, Decompress, FlushCompress, FlushDecompress,
        Status,
    };
    use std::io::Write;

    #[cfg(not(feature = "zlib"))]
    mod constants {
        pub const Z_NO_COMPRESSION: i32 = 0;
        pub const Z_BEST_COMPRESSION: i32 = 9;
        pub const Z_BEST_SPEED: i32 = 1;
        pub const Z_DEFAULT_COMPRESSION: i32 = -1;
        pub const Z_NO_FLUSH: i32 = 0;
        pub const Z_PARTIAL_FLUSH: i32 = 1;
        pub const Z_SYNC_FLUSH: i32 = 2;
        pub const Z_FULL_FLUSH: i32 = 3;
        // not sure what the value here means, but it's the only compression method zlibmodule
        // supports, so it doesn't really matter
        pub const Z_DEFLATED: i32 = 8;
    }
    #[cfg(feature = "zlib")]
    use libz_sys as constants;

    #[pyattr]
    use constants::{
        Z_BEST_COMPRESSION, Z_BEST_SPEED, Z_DEFAULT_COMPRESSION, Z_DEFLATED as DEFLATED,
        Z_FULL_FLUSH, Z_NO_COMPRESSION, Z_NO_FLUSH, Z_PARTIAL_FLUSH, Z_SYNC_FLUSH,
    };

    #[cfg(feature = "zlib")]
    #[pyattr]
    use libz_sys::{
        Z_BLOCK, Z_DEFAULT_STRATEGY, Z_FILTERED, Z_FINISH, Z_FIXED, Z_HUFFMAN_ONLY, Z_RLE, Z_TREES,
    };

    // copied from zlibmodule.c (commit 530f506ac91338)
    #[pyattr]
    const MAX_WBITS: i8 = 15;
    #[pyattr]
    const DEF_BUF_SIZE: usize = 16 * 1024;
    #[pyattr]
    const DEF_MEM_LEVEL: u8 = 8;

    #[pyattr(once)]
    fn error(vm: &VirtualMachine) -> PyTypeRef {
        vm.ctx.new_exception_type(
            "zlib",
            "error",
            Some(vec![vm.ctx.exceptions.exception_type.to_owned()]),
        )
    }

    #[pyfunction]
    fn adler32(data: ArgBytesLike, begin_state: OptionalArg<PyIntRef>) -> u32 {
        data.with_ref(|data| {
            let begin_state = begin_state.map_or(1, |i| i.as_u32_mask());

            let mut hasher = Adler32::from_value(begin_state);
            hasher.update_buffer(data);
            hasher.hash()
        })
    }

    #[pyfunction]
    fn crc32(data: ArgBytesLike, begin_state: OptionalArg<PyIntRef>) -> u32 {
        crate::binascii::crc32(data, begin_state)
    }

    #[derive(FromArgs)]
    struct PyFuncCompressArgs {
        #[pyarg(positional)]
        data: ArgBytesLike,
        #[pyarg(any, default = "Level::new(Z_DEFAULT_COMPRESSION)")]
        level: Level,
        #[pyarg(any, default = "ArgPrimitiveIndex { value: MAX_WBITS }")]
        wbits: ArgPrimitiveIndex<i8>,
    }

    /// Returns a bytes object containing compressed data.
    #[pyfunction]
    fn compress(args: PyFuncCompressArgs, vm: &VirtualMachine) -> PyResult<PyBytesRef> {
        let PyFuncCompressArgs {
            data,
            level,
            ref wbits,
        } = args;
        let level = level.ok_or_else(|| new_zlib_error("Bad compression level", vm))?;

        let encoded_bytes = if args.wbits.value == MAX_WBITS {
            let mut encoder = ZlibEncoder::new(Vec::new(), level);
            data.with_ref(|input_bytes| encoder.write_all(input_bytes).unwrap());
            encoder.finish().unwrap()
        } else {
            let mut inner = CompressInner::new(InitOptions::new(wbits.value, vm)?.compress(level));
            data.with_ref(|input_bytes| inner.compress(input_bytes, vm))?;
            inner.flush(vm)?
        };
        Ok(vm.ctx.new_bytes(encoded_bytes))
    }

    enum InitOptions {
        Standard {
            header: bool,
            // [De]Compress::new_with_window_bits is only enabled for zlib; miniz_oxide doesn't
            // support wbits (yet?)
            #[cfg(feature = "zlib")]
            wbits: u8,
        },
        #[cfg(feature = "zlib")]
        Gzip { wbits: u8 },
    }

    impl InitOptions {
        fn new(wbits: i8, vm: &VirtualMachine) -> PyResult<InitOptions> {
            let header = wbits > 0;
            let wbits = wbits.unsigned_abs();
            match wbits {
                9..=15 => Ok(InitOptions::Standard {
                    header,
                    #[cfg(feature = "zlib")]
                    wbits,
                }),
                #[cfg(feature = "zlib")]
                25..=31 => Ok(InitOptions::Gzip { wbits: wbits - 16 }),
                _ => Err(vm.new_value_error("Invalid initialization option".to_owned())),
            }
        }

        fn decompress(self) -> Decompress {
            match self {
                #[cfg(not(feature = "zlib"))]
                Self::Standard { header } => Decompress::new(header),
                #[cfg(feature = "zlib")]
                Self::Standard { header, wbits } => Decompress::new_with_window_bits(header, wbits),
                #[cfg(feature = "zlib")]
                Self::Gzip { wbits } => Decompress::new_gzip(wbits),
            }
        }
        fn compress(self, level: Compression) -> Compress {
            match self {
                #[cfg(not(feature = "zlib"))]
                Self::Standard { header } => Compress::new(level, header),
                #[cfg(feature = "zlib")]
                Self::Standard { header, wbits } => {
                    Compress::new_with_window_bits(level, header, wbits)
                }
                #[cfg(feature = "zlib")]
                Self::Gzip { wbits } => Compress::new_gzip(level, wbits),
            }
        }
    }

    fn _decompress(
        mut data: &[u8],
        d: &mut Decompress,
        bufsize: usize,
        max_length: Option<usize>,
        is_flush: bool,
        vm: &VirtualMachine,
    ) -> PyResult<(Vec<u8>, bool)> {
        if data.is_empty() {
            return Ok((Vec::new(), true));
        }
        let mut buf = Vec::new();

        loop {
            let final_chunk = data.len() <= CHUNKSIZE;
            let chunk = if final_chunk {
                data
            } else {
                &data[..CHUNKSIZE]
            };
            // if this is the final chunk, finish it
            let flush = if is_flush {
                if final_chunk {
                    FlushDecompress::Finish
                } else {
                    FlushDecompress::None
                }
            } else {
                FlushDecompress::Sync
            };
            loop {
                let additional = if let Some(max_length) = max_length {
                    std::cmp::min(bufsize, max_length - buf.capacity())
                } else {
                    bufsize
                };
                if additional == 0 {
                    return Ok((buf, false));
                }

                buf.reserve_exact(additional);
                let prev_in = d.total_in();
                let status = d
                    .decompress_vec(chunk, &mut buf, flush)
                    .map_err(|_| new_zlib_error("invalid input data", vm))?;
                let consumed = d.total_in() - prev_in;
                data = &data[consumed as usize..];
                let stream_end = status == Status::StreamEnd;
                if stream_end || data.is_empty() {
                    // we've reached the end of the stream, we're done
                    buf.shrink_to_fit();
                    return Ok((buf, stream_end));
                } else if !chunk.is_empty() && consumed == 0 {
                    // we're gonna need a bigger buffer
                    continue;
                } else {
                    // next chunk
                    break;
                }
            }
        }
    }

    #[derive(FromArgs)]
    struct PyFuncDecompressArgs {
        #[pyarg(positional)]
        data: ArgBytesLike,
        #[pyarg(any, default = "ArgPrimitiveIndex { value: MAX_WBITS }")]
        wbits: ArgPrimitiveIndex<i8>,
        #[pyarg(any, default = "ArgPrimitiveIndex { value: DEF_BUF_SIZE }")]
        bufsize: ArgPrimitiveIndex<usize>,
    }

    /// Returns a bytes object containing the uncompressed data.
    #[pyfunction]
    fn decompress(args: PyFuncDecompressArgs, vm: &VirtualMachine) -> PyResult<Vec<u8>> {
        let PyFuncDecompressArgs {
            data,
            wbits,
            bufsize,
        } = args;
        data.with_ref(|data| {
            let mut d = InitOptions::new(wbits.value, vm)?.decompress();
            let (buf, stream_end) = _decompress(data, &mut d, bufsize.value, None, false, vm)?;
            if !stream_end {
                return Err(new_zlib_error(
                    "Error -5 while decompressing data: incomplete or truncated stream",
                    vm,
                ));
            }
            Ok(buf)
        })
    }

    #[derive(FromArgs)]
    struct DecompressobjArgs {
        #[pyarg(any, default = "ArgPrimitiveIndex { value: MAX_WBITS }")]
        wbits: ArgPrimitiveIndex<i8>,
        #[cfg(feature = "zlib")]
        #[pyarg(any, optional)]
        _zdict: OptionalArg<ArgBytesLike>,
    }

    #[pyfunction]
    fn decompressobj(args: DecompressobjArgs, vm: &VirtualMachine) -> PyResult<PyDecompress> {
        #[allow(unused_mut)]
        let mut decompress = InitOptions::new(args.wbits.value, vm)?.decompress();
        #[cfg(feature = "zlib")]
        if let OptionalArg::Present(_dict) = args._zdict {
            // FIXME: always fails
            // dict.with_ref(|d| decompress.set_dictionary(d));
        }
        Ok(PyDecompress {
            decompress: PyMutex::new(decompress),
            eof: AtomicCell::new(false),
            unused_data: PyMutex::new(PyBytes::from(vec![]).into_ref(&vm.ctx)),
            unconsumed_tail: PyMutex::new(PyBytes::from(vec![]).into_ref(&vm.ctx)),
        })
    }
    #[pyattr]
    #[pyclass(name = "Decompress")]
    #[derive(Debug, PyPayload)]
    struct PyDecompress {
        decompress: PyMutex<Decompress>,
        eof: AtomicCell<bool>,
        unused_data: PyMutex<PyBytesRef>,
        unconsumed_tail: PyMutex<PyBytesRef>,
    }
    #[pyclass]
    impl PyDecompress {
        #[pygetset]
        fn eof(&self) -> bool {
            self.eof.load()
        }
        #[pygetset]
        fn unused_data(&self) -> PyBytesRef {
            self.unused_data.lock().clone()
        }
        #[pygetset]
        fn unconsumed_tail(&self) -> PyBytesRef {
            self.unconsumed_tail.lock().clone()
        }

        fn save_unused_input(
            &self,
            d: &Decompress,
            data: &[u8],
            stream_end: bool,
            orig_in: u64,
            vm: &VirtualMachine,
        ) {
            let leftover = &data[(d.total_in() - orig_in) as usize..];

            if stream_end && !leftover.is_empty() {
                let mut unused_data = self.unused_data.lock();
                let unused: Vec<_> = unused_data
                    .as_bytes()
                    .iter()
                    .chain(leftover)
                    .copied()
                    .collect();
                *unused_data = vm.ctx.new_pyref(unused);
            }
        }

        #[pymethod]
        fn decompress(&self, args: DecompressArgs, vm: &VirtualMachine) -> PyResult<Vec<u8>> {
            let max_length = args.max_length.value;
            let max_length = (max_length != 0).then_some(max_length);
            let data = args.data.borrow_buf();
            let data = &*data;

            let mut d = self.decompress.lock();
            let orig_in = d.total_in();

            let (ret, stream_end) =
                match _decompress(data, &mut d, DEF_BUF_SIZE, max_length, false, vm) {
                    Ok((buf, true)) => {
                        self.eof.store(true);
                        (Ok(buf), true)
                    }
                    Ok((buf, false)) => (Ok(buf), false),
                    Err(err) => (Err(err), false),
                };
            self.save_unused_input(&d, data, stream_end, orig_in, vm);

            let leftover = if stream_end {
                b""
            } else {
                &data[(d.total_in() - orig_in) as usize..]
            };

            let mut unconsumed_tail = self.unconsumed_tail.lock();
            if !leftover.is_empty() || !unconsumed_tail.is_empty() {
                *unconsumed_tail = PyBytes::from(leftover.to_owned()).into_ref(&vm.ctx);
            }

            ret
        }

        #[pymethod]
        fn flush(&self, length: OptionalArg<ArgSize>, vm: &VirtualMachine) -> PyResult<Vec<u8>> {
            let length = match length {
                OptionalArg::Present(l) => {
                    let l: isize = l.into();
                    if l <= 0 {
                        return Err(
                            vm.new_value_error("length must be greater than zero".to_owned())
                        );
                    } else {
                        l as usize
                    }
                }
                OptionalArg::Missing => DEF_BUF_SIZE,
            };

            let mut data = self.unconsumed_tail.lock();
            let mut d = self.decompress.lock();

            let orig_in = d.total_in();

            let (ret, stream_end) = match _decompress(&data, &mut d, length, None, true, vm) {
                Ok((buf, stream_end)) => (Ok(buf), stream_end),
                Err(err) => (Err(err), false),
            };
            self.save_unused_input(&d, &data, stream_end, orig_in, vm);

            *data = PyBytes::from(Vec::new()).into_ref(&vm.ctx);

            // TODO: drop the inner decompressor, somehow
            // if stream_end {
            //
            // }
            ret
        }
    }

    #[derive(FromArgs)]
    struct DecompressArgs {
        #[pyarg(positional)]
        data: ArgBytesLike,
        #[pyarg(
            any,
            default = "rustpython_vm::function::ArgPrimitiveIndex { value: 0 }"
        )]
        max_length: ArgPrimitiveIndex<usize>,
    }

    #[derive(FromArgs)]
    #[allow(dead_code)] // FIXME: use args
    struct CompressobjArgs {
        #[pyarg(any, default = "Level::new(Z_DEFAULT_COMPRESSION)")]
        level: Level,
        // only DEFLATED is valid right now, it's w/e
        #[pyarg(any, default = "DEFLATED")]
        _method: i32,
        #[pyarg(any, default = "ArgPrimitiveIndex { value: MAX_WBITS }")]
        wbits: ArgPrimitiveIndex<i8>,
        #[pyarg(any, name = "_memLevel", default = "DEF_MEM_LEVEL")]
        _mem_level: u8,
        #[cfg(feature = "zlib")]
        #[pyarg(any, default = "Z_DEFAULT_STRATEGY")]
        _strategy: i32,
        #[cfg(feature = "zlib")]
        #[pyarg(any, optional)]
        zdict: Option<ArgBytesLike>,
    }

    #[pyfunction]
    fn compressobj(args: CompressobjArgs, vm: &VirtualMachine) -> PyResult<PyCompress> {
        let CompressobjArgs {
            level,
            wbits,
            #[cfg(feature = "zlib")]
            zdict,
            ..
        } = args;
        let level =
            level.ok_or_else(|| vm.new_value_error("invalid initialization option".to_owned()))?;
        #[allow(unused_mut)]
        let mut compress = InitOptions::new(wbits.value, vm)?.compress(level);
        #[cfg(feature = "zlib")]
        if let Some(zdict) = zdict {
            zdict.with_ref(|zdict| compress.set_dictionary(zdict).unwrap());
        }
        Ok(PyCompress {
            inner: PyMutex::new(CompressInner::new(compress)),
        })
    }

    #[derive(Debug)]
    struct CompressInner {
        compress: Compress,
        unconsumed: Vec<u8>,
    }

    #[pyattr]
    #[pyclass(name = "Compress")]
    #[derive(Debug, PyPayload)]
    struct PyCompress {
        inner: PyMutex<CompressInner>,
    }

    #[pyclass]
    impl PyCompress {
        #[pymethod]
        fn compress(&self, data: ArgBytesLike, vm: &VirtualMachine) -> PyResult<Vec<u8>> {
            let mut inner = self.inner.lock();
            data.with_ref(|b| inner.compress(b, vm))
        }

        // TODO: mode argument isn't used
        #[pymethod]
        fn flush(&self, _mode: OptionalArg<i32>, vm: &VirtualMachine) -> PyResult<Vec<u8>> {
            self.inner.lock().flush(vm)
        }

        // TODO: This is an optional feature of Compress
        // #[pymethod]
        // #[pymethod(magic)]
        // #[pymethod(name = "__deepcopy__")]
        // fn copy(&self) -> Self {
        //     todo!("<flate2::Compress as Clone>")
        // }
    }

    const CHUNKSIZE: usize = u32::MAX as usize;

    impl CompressInner {
        fn new(compress: Compress) -> Self {
            Self {
                compress,
                unconsumed: Vec::new(),
            }
        }
        fn compress(&mut self, data: &[u8], vm: &VirtualMachine) -> PyResult<Vec<u8>> {
            let orig_in = self.compress.total_in() as usize;
            let mut cur_in = 0;
            let unconsumed = std::mem::take(&mut self.unconsumed);
            let mut buf = Vec::new();

            'outer: for chunk in unconsumed.chunks(CHUNKSIZE).chain(data.chunks(CHUNKSIZE)) {
                while cur_in < chunk.len() {
                    buf.reserve(DEF_BUF_SIZE);
                    let status = self
                        .compress
                        .compress_vec(&chunk[cur_in..], &mut buf, FlushCompress::None)
                        .map_err(|_| {
                            self.unconsumed.extend_from_slice(&data[cur_in..]);
                            new_zlib_error("error while compressing", vm)
                        })?;
                    cur_in = (self.compress.total_in() as usize) - orig_in;
                    match status {
                        Status::Ok => continue,
                        Status::StreamEnd => break 'outer,
                        _ => break,
                    }
                }
            }
            self.unconsumed.extend_from_slice(&data[cur_in..]);

            buf.shrink_to_fit();
            Ok(buf)
        }

        // TODO: flush mode (FlushDecompress) parameter
        fn flush(&mut self, vm: &VirtualMachine) -> PyResult<Vec<u8>> {
            let data = std::mem::take(&mut self.unconsumed);
            let mut data_it = data.chunks(CHUNKSIZE);
            let mut buf = Vec::new();

            loop {
                let chunk = data_it.next().unwrap_or(&[]);
                if buf.len() == buf.capacity() {
                    buf.reserve(DEF_BUF_SIZE);
                }
                let status = self
                    .compress
                    .compress_vec(chunk, &mut buf, FlushCompress::Finish)
                    .map_err(|_| new_zlib_error("error while compressing", vm))?;
                match status {
                    Status::StreamEnd => break,
                    _ => continue,
                }
            }

            buf.shrink_to_fit();
            Ok(buf)
        }
    }

    fn new_zlib_error(message: &str, vm: &VirtualMachine) -> PyBaseExceptionRef {
        vm.new_exception_msg(vm.class("zlib", "error"), message.to_owned())
    }

    struct Level(Option<flate2::Compression>);

    impl Level {
        fn new(level: i32) -> Self {
            let compression = match level {
                Z_DEFAULT_COMPRESSION => Compression::default(),
                valid_level @ Z_NO_COMPRESSION..=Z_BEST_COMPRESSION => {
                    Compression::new(valid_level as u32)
                }
                _ => return Self(None),
            };
            Self(Some(compression))
        }
        fn ok_or_else(
            self,
            f: impl FnOnce() -> PyBaseExceptionRef,
        ) -> PyResult<flate2::Compression> {
            self.0.ok_or_else(f)
        }
    }

    impl<'a> TryFromBorrowedObject<'a> for Level {
        fn try_from_borrowed_object(vm: &VirtualMachine, obj: &'a PyObject) -> PyResult<Self> {
            let int: i32 = obj.try_index(vm)?.try_to_primitive(vm)?;
            Ok(Self::new(int))
        }
    }
}
