use crate::driver::FileDriver;
use crate::utils::{IoCursor, ReadUtils};
use crate::zip::encryption::{AesDecoder, Keys, WeakDecoder};
use crate::zip::structs::FILE_HEADER_SIGNATURE;
use crate::zip::{CompressionMethod, EncryptionMethod, ZipError, ZipFileInfo, ZipResult};
use aes::cipher::KeyInit;
use aes::{Aes128, Aes192, Aes256};
use bzip2::read::BzDecoder;
use flate2::read::DeflateDecoder;
use liblzma::read::XzDecoder;
use liblzma::stream::{Filters, LzmaOptions, Stream};
use pbkdf2::pbkdf2_hmac_array;
use sha1::Sha1;
use std::io::{
    BufReader, Error as IoError, ErrorKind as IoErrorKind, Read, Result as IoResult, Seek, SeekFrom,
};
use zstd::stream::Decoder as ZstdDecoder;

enum Encryption<Io: Read> {
    None(Io),
    Weak(WeakDecoder<Io>),
    Aes128(AesDecoder<Io, Aes128>),
    Aes192(AesDecoder<Io, Aes192>),
    Aes256(AesDecoder<Io, Aes256>),
}

impl<Io: Read + Seek> Encryption<IoCursor<Io>> {
    #[inline]
    pub fn new(
        mut io: Io,
        cursor: u64,
        end: u64,
        info: &ZipFileInfo,
        password: Option<&[u8]>,
    ) -> ZipResult<Self> {
        Ok(match info.encryption_method {
            EncryptionMethod::None => Self::None(IoCursor::new(io, cursor, end)),
            EncryptionMethod::Weak => {
                let mut keys = Keys::new();

                for b in password.ok_or(ZipError::PasswordIsNotSpecified)? {
                    keys.update(*b)
                }

                let header = io.read_arr::<12>()?;
                for b in &header[..11] {
                    keys.decode_byte(*b);
                }

                if keys.decode_byte(header[11]) != info.password_check() {
                    return Err(ZipError::WrongPassword);
                }

                Self::Weak(WeakDecoder::new(IoCursor::new(io, cursor + 12, end), keys))
            }
            EncryptionMethod::Aes128 => {
                let header = io.read_arr::<10>()?;
                let salt = &header[..8];
                let check = &header[8..];

                let hash = pbkdf2_hmac_array::<Sha1, 34>(
                    password.ok_or(ZipError::PasswordIsNotSpecified)?,
                    salt,
                    1000,
                );
                let key = &hash[..16];

                if check != &hash[32..] {
                    return Err(ZipError::WrongPassword);
                }

                Self::Aes128(AesDecoder::new(
                    IoCursor::new(io, cursor + 10, end - 10),
                    Aes128::new(key.into()),
                )?)
            }
            EncryptionMethod::Aes192 => {
                let header = io.read_arr::<14>()?;
                let salt = &header[..12];
                let check = &header[12..];

                let hash = pbkdf2_hmac_array::<Sha1, 50>(
                    password.ok_or(ZipError::PasswordIsNotSpecified)?,
                    salt,
                    1000,
                );
                let key = &hash[..24];

                if check != &hash[48..] {
                    return Err(ZipError::WrongPassword);
                }

                Self::Aes192(AesDecoder::new(
                    IoCursor::new(io, cursor + 14, end - 10),
                    Aes192::new(key.into()),
                )?)
            }
            EncryptionMethod::Aes256 => {
                let header = io.read_arr::<18>()?;
                let salt = &header[..16];
                let check = &header[16..];

                let hash = pbkdf2_hmac_array::<Sha1, 66>(
                    password.ok_or(ZipError::PasswordIsNotSpecified)?,
                    salt,
                    1000,
                );
                let key = &hash[..32];

                if check != &hash[64..] {
                    return Err(ZipError::WrongPassword);
                }

                Self::Aes256(AesDecoder::new(
                    IoCursor::new(io, cursor + 18, end - 10),
                    Aes256::new(key.into()),
                )?)
            }
            EncryptionMethod::Unsupported => {
                return Err(ZipError::Unsupported("encryption method"))
            }
        })
    }
}

impl<Io: Read> Read for Encryption<Io> {
    #[inline]
    fn read(&mut self, buf: &mut [u8]) -> IoResult<usize> {
        match self {
            Self::None(io) => io.read(buf),
            Self::Weak(io) => io.read(buf),
            Self::Aes128(io) => io.read(buf),
            Self::Aes192(io) => io.read(buf),
            Self::Aes256(io) => io.read(buf),
        }
    }
}

impl<Io: Read + Seek> Seek for Encryption<Io> {
    #[inline]
    fn seek(&mut self, pos: SeekFrom) -> IoResult<u64> {
        match self {
            Self::None(io) => io.seek(pos),
            _ => Err(IoError::new(
                IoErrorKind::Unsupported,
                ZipError::UnseekableFile,
            )),
        }
    }
}

enum Compression<Io: Read> {
    Store(Io),
    Deflate(DeflateDecoder<Io>),
    BZip2(BzDecoder<Io>),
    Zstd(ZstdDecoder<'static, BufReader<Io>>),
    Xz(XzDecoder<Io>),
}

impl<Io: Read + Seek> Compression<Io> {
    #[inline]
    pub fn new(mut io: Io, info: &ZipFileInfo) -> ZipResult<Self> {
        Ok(match info.compression_method {
            CompressionMethod::Store => Self::Store(io),
            CompressionMethod::Deflate => Self::Deflate(DeflateDecoder::new(io)),
            CompressionMethod::BZip2 => Self::BZip2(BzDecoder::new(io)),
            CompressionMethod::Lzma => {
                let buf = io.read_arr::<9>()?;
                Compression::Xz(XzDecoder::new_stream(
                    io,
                    Stream::new_raw_decoder(
                        Filters::new().lzma1(
                            LzmaOptions::new()
                                .literal_context_bits((buf[4] % 9) as u32)
                                .literal_position_bits((buf[4] / 9 % 5) as u32)
                                .position_bits((buf[4] / 45) as u32)
                                .dict_size(
                                    u32::from_le_bytes(buf[5..9].try_into().unwrap()).max(4096),
                                ),
                        ),
                    )
                    .unwrap(),
                ))
            }
            CompressionMethod::Zstd => Self::Zstd(ZstdDecoder::new(io)?),
            CompressionMethod::Xz => Self::Xz(XzDecoder::new(io)),
            CompressionMethod::Unsupported => {
                return Err(ZipError::Unsupported("compression method"));
            }
        })
    }
}

impl<Io: Read> Read for Compression<Io> {
    #[inline]
    fn read(&mut self, buf: &mut [u8]) -> IoResult<usize> {
        match self {
            Compression::Store(io) => io.read(buf),
            Compression::Deflate(io) => io.read(buf),
            Compression::BZip2(io) => io.read(buf),
            Compression::Zstd(io) => io.read(buf),
            Compression::Xz(io) => io.read(buf),
        }
    }
}

impl<Io: Read + Seek> Seek for Compression<Io> {
    #[inline]
    fn seek(&mut self, pos: SeekFrom) -> IoResult<u64> {
        match self {
            Compression::Store(io) => io.seek(pos),
            _ => Err(IoError::new(
                IoErrorKind::Unsupported,
                ZipError::UnseekableFile,
            )),
        }
    }
}

pub struct ZipFileReader<'d, Io: Read> {
    io: Compression<Encryption<IoCursor<&'d mut Io>>>,
    info: &'d ZipFileInfo,
}

impl<'d, Io: Read> FileDriver for ZipFileReader<'d, Io> {
    type Io = Io;
    type FileInfo = ZipFileInfo;
}

impl<'d, Io: Read + Seek> ZipFileReader<'d, Io> {
    pub(crate) fn new(
        io: &'d mut Io,
        info: &'d ZipFileInfo,
        password: Option<&[u8]>,
    ) -> ZipResult<Self> {
        io.seek(SeekFrom::Start(info.header_pointer))?;

        let buf = io.read_arr::<30>()?;
        if buf[..4] != FILE_HEADER_SIGNATURE {
            return Err(ZipError::InvalidSignature("FileHeader"));
        }

        let cursor = io.seek(SeekFrom::Start(
            info.header_pointer
                + 30
                + u16::from_le_bytes(buf[26..28].try_into().unwrap()) as u64
                + u16::from_le_bytes(buf[28..30].try_into().unwrap()) as u64,
        ))?;
        let end = cursor + info.compressed_size;

        Ok(Self {
            io: Compression::new(Encryption::new(io, cursor, end, info, password)?, info)?,
            info,
        })
    }

    pub fn info(&self) -> &ZipFileInfo {
        self.info
    }

    pub fn is_seekable(&self) -> bool {
        match self.io {
            Compression::Store(Encryption::None(..)) => true,
            _ => false,
        }
    }
}

impl<'d, Io: Read> Read for ZipFileReader<'d, Io> {
    fn read(&mut self, buf: &mut [u8]) -> IoResult<usize> {
        self.io.read(buf)
    }
}

impl<'d, Io: Read + Seek> Seek for ZipFileReader<'d, Io> {
    fn seek(&mut self, pos: SeekFrom) -> IoResult<u64> {
        self.io.seek(pos)
    }
}