use async_compression::tokio::write::{BrotliEncoder, GzipEncoder, ZstdEncoder};
use bytes::Bytes;
use futures::stream::{Stream, StreamExt};
use tokio::io::AsyncWriteExt;
use tracing::error;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CompressionEncoding {
    Zstd,
    Brotli,
    Gzip,
    Identity,
}

impl CompressionEncoding {
    pub fn from_accept_encoding(accept_encoding: Option<&str>) -> Self {
        let accept = accept_encoding.unwrap_or("");

        if accept.contains("zstd") {
            Self::Zstd
        } else if accept.contains("br") {
            Self::Brotli
        } else if accept.contains("gzip") {
            Self::Gzip
        } else {
            Self::Identity
        }
    }

    pub fn as_header_value(&self) -> Option<&'static str> {
        match self {
            Self::Zstd => Some("zstd"),
            Self::Brotli => Some("br"),
            Self::Gzip => Some("gzip"),
            Self::Identity => None,
        }
    }
}

pub fn compress_stream<S>(
    input: S,
    encoding: CompressionEncoding,
) -> impl Stream<Item = Result<Bytes, std::io::Error>>
where
    S: Stream<Item = Result<Bytes, std::io::Error>> + Send + 'static,
{
    match encoding {
        CompressionEncoding::Zstd => {
            let pinned_input = Box::pin(input);
            futures::stream::unfold(
                (pinned_input, Some(ZstdCompressor::new()), false),
                |(mut input, mut compressor, finished)| async move {
                    if finished {
                        return None;
                    }

                    match input.next().await {
                        Some(Ok(chunk)) => {
                            let mut comp = compressor.take()?;
                            match comp.compress_and_flush(&chunk).await {
                                Ok(compressed_chunk) => {
                                    Some((Ok(compressed_chunk), (input, Some(comp), false)))
                                }
                                Err(e) => {
                                    error!("Zstd compression error: {}", e);
                                    Some((Err(e), (input, None, true)))
                                }
                            }
                        }
                        Some(Err(e)) => Some((Err(e), (input, None, true))),
                        None => {
                            if let Some(comp) = compressor.take() {
                                match comp.finish().await {
                                    Ok(final_chunk) => Some((Ok(final_chunk), (input, None, true))),
                                    Err(e) => {
                                        error!("Zstd finalization error: {}", e);
                                        Some((Err(e), (input, None, true)))
                                    }
                                }
                            } else {
                                None
                            }
                        }
                    }
                },
            )
            .boxed()
        }
        CompressionEncoding::Brotli => {
            let pinned_input = Box::pin(input);
            futures::stream::unfold(
                (pinned_input, Some(BrotliCompressor::new()), false),
                |(mut input, mut compressor, finished)| async move {
                    if finished {
                        return None;
                    }

                    match input.next().await {
                        Some(Ok(chunk)) => {
                            let mut comp = compressor.take()?;
                            match comp.compress_and_flush(&chunk).await {
                                Ok(compressed_chunk) => {
                                    Some((Ok(compressed_chunk), (input, Some(comp), false)))
                                }
                                Err(e) => {
                                    error!("Brotli compression error: {}", e);
                                    Some((Err(e), (input, None, true)))
                                }
                            }
                        }
                        Some(Err(e)) => Some((Err(e), (input, None, true))),
                        None => {
                            if let Some(comp) = compressor.take() {
                                match comp.finish().await {
                                    Ok(final_chunk) => Some((Ok(final_chunk), (input, None, true))),
                                    Err(e) => {
                                        error!("Brotli finalization error: {}", e);
                                        Some((Err(e), (input, None, true)))
                                    }
                                }
                            } else {
                                None
                            }
                        }
                    }
                },
            )
            .boxed()
        }
        CompressionEncoding::Gzip => {
            let pinned_input = Box::pin(input);
            futures::stream::unfold(
                (pinned_input, Some(GzipCompressor::new()), false),
                |(mut input, mut compressor, finished)| async move {
                    if finished {
                        return None;
                    }

                    match input.next().await {
                        Some(Ok(chunk)) => {
                            let mut comp = compressor.take()?;
                            match comp.compress_and_flush(&chunk).await {
                                Ok(compressed_chunk) => {
                                    Some((Ok(compressed_chunk), (input, Some(comp), false)))
                                }
                                Err(e) => {
                                    error!("Gzip compression error: {}", e);
                                    Some((Err(e), (input, None, true)))
                                }
                            }
                        }
                        Some(Err(e)) => Some((Err(e), (input, None, true))),
                        None => {
                            if let Some(comp) = compressor.take() {
                                match comp.finish().await {
                                    Ok(final_chunk) => Some((Ok(final_chunk), (input, None, true))),
                                    Err(e) => {
                                        error!("Gzip finalization error: {}", e);
                                        Some((Err(e), (input, None, true)))
                                    }
                                }
                            } else {
                                None
                            }
                        }
                    }
                },
            )
            .boxed()
        }
        CompressionEncoding::Identity => Box::pin(input).boxed(),
    }
}

struct ZstdCompressor {
    encoder: ZstdEncoder<Vec<u8>>,
}

impl ZstdCompressor {
    fn new() -> Self {
        Self { encoder: ZstdEncoder::new(Vec::new()) }
    }

    async fn compress_and_flush(&mut self, data: &[u8]) -> std::io::Result<Bytes> {
        self.encoder.write_all(data).await?;
        self.encoder.flush().await?;

        let inner = self.encoder.get_mut();
        let compressed = std::mem::take(inner);

        Ok(Bytes::from(compressed))
    }

    async fn finish(mut self) -> std::io::Result<Bytes> {
        self.encoder.shutdown().await?;
        let final_bytes = self.encoder.into_inner();
        Ok(Bytes::from(final_bytes))
    }
}

struct BrotliCompressor {
    encoder: BrotliEncoder<Vec<u8>>,
}

impl BrotliCompressor {
    fn new() -> Self {
        Self { encoder: BrotliEncoder::new(Vec::new()) }
    }

    async fn compress_and_flush(&mut self, data: &[u8]) -> std::io::Result<Bytes> {
        self.encoder.write_all(data).await?;
        self.encoder.flush().await?;

        let inner = self.encoder.get_mut();
        let compressed = std::mem::take(inner);

        Ok(Bytes::from(compressed))
    }

    async fn finish(mut self) -> std::io::Result<Bytes> {
        self.encoder.shutdown().await?;
        let final_bytes = self.encoder.into_inner();
        Ok(Bytes::from(final_bytes))
    }
}

struct GzipCompressor {
    encoder: GzipEncoder<Vec<u8>>,
}

impl GzipCompressor {
    fn new() -> Self {
        Self { encoder: GzipEncoder::new(Vec::new()) }
    }

    async fn compress_and_flush(&mut self, data: &[u8]) -> std::io::Result<Bytes> {
        self.encoder.write_all(data).await?;
        self.encoder.flush().await?;

        let inner = self.encoder.get_mut();
        let compressed = std::mem::take(inner);

        Ok(Bytes::from(compressed))
    }

    async fn finish(mut self) -> std::io::Result<Bytes> {
        self.encoder.shutdown().await?;
        let final_bytes = self.encoder.into_inner();
        Ok(Bytes::from(final_bytes))
    }
}

#[cfg(test)]
#[allow(clippy::disallowed_methods)]
mod tests {
    use super::*;
    use futures::stream;

    #[tokio::test]
    async fn test_gzip_compression() {
        let data = vec![Ok(Bytes::from("Hello ")), Ok(Bytes::from("World")), Ok(Bytes::from("!"))];

        let input_stream = stream::iter(data);
        let mut compressed_stream = compress_stream(input_stream, CompressionEncoding::Gzip);

        let mut chunks = Vec::new();
        while let Some(chunk) = compressed_stream.next().await {
            chunks.push(chunk.unwrap());
        }

        assert!(!chunks.is_empty());
        assert!(chunks.iter().all(|c| !c.is_empty()));
    }

    #[tokio::test]
    async fn test_brotli_compression() {
        let data = vec![Ok(Bytes::from("Hello ")), Ok(Bytes::from("World"))];

        let input_stream = stream::iter(data);
        let mut compressed_stream = compress_stream(input_stream, CompressionEncoding::Brotli);

        let mut chunks = Vec::new();
        while let Some(chunk) = compressed_stream.next().await {
            chunks.push(chunk.unwrap());
        }

        assert!(!chunks.is_empty());
    }

    #[tokio::test]
    async fn test_zstd_compression() {
        let data = vec![Ok(Bytes::from("Hello ")), Ok(Bytes::from("World")), Ok(Bytes::from("!"))];

        let input_stream = stream::iter(data);
        let mut compressed_stream = compress_stream(input_stream, CompressionEncoding::Zstd);

        let mut chunks = Vec::new();
        while let Some(chunk) = compressed_stream.next().await {
            chunks.push(chunk.unwrap());
        }

        assert!(!chunks.is_empty());
        assert!(chunks.iter().all(|c| !c.is_empty()));
    }

    #[tokio::test]
    async fn test_identity_passthrough() {
        let input_stream = stream::iter(vec![Ok(Bytes::from("Hello ")), Ok(Bytes::from("World"))]);
        let mut compressed_stream = compress_stream(input_stream, CompressionEncoding::Identity);

        let mut chunks = Vec::new();
        while let Some(chunk) = compressed_stream.next().await {
            chunks.push(chunk.unwrap());
        }

        assert_eq!(chunks.len(), 2);
        assert_eq!(chunks[0], Bytes::from("Hello "));
        assert_eq!(chunks[1], Bytes::from("World"));
    }

    #[test]
    fn test_encoding_from_accept_header() {
        assert_eq!(
            CompressionEncoding::from_accept_encoding(Some("gzip, deflate, br, zstd")),
            CompressionEncoding::Zstd
        );
        assert_eq!(
            CompressionEncoding::from_accept_encoding(Some("gzip, deflate, br")),
            CompressionEncoding::Brotli
        );
        assert_eq!(
            CompressionEncoding::from_accept_encoding(Some("gzip, deflate")),
            CompressionEncoding::Gzip
        );
        assert_eq!(
            CompressionEncoding::from_accept_encoding(Some("identity")),
            CompressionEncoding::Identity
        );
        assert_eq!(CompressionEncoding::from_accept_encoding(None), CompressionEncoding::Identity);
    }
}
