use flate2::read::{GzDecoder, ZlibDecoder};
use std::fs::File;
use std::io::{BufRead, BufReader, Read};
use std::path::Path;

use crate::error::{BunbunError, Result};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Codec {
    None,
    Gzip,
    Bgzf,
    Zlib,
    Zstd,
    Bzip2,
}

impl Codec {
    pub fn from_extension(path: &Path) -> Self {
        let name = path.file_name().and_then(|n| n.to_str()).unwrap_or("");
        if name.ends_with(".gz") || name.ends_with(".bgz") {
            Self::Gzip
        } else if name.ends_with(".zst") {
            Self::Zstd
        } else if name.ends_with(".bz2") {
            Self::Bzip2
        } else {
            Self::None
        }
    }

    pub fn from_magic(bytes: &[u8]) -> Self {
        if bytes.len() >= 2 && bytes[0] == 0x1f && bytes[1] == 0x8b {
            Self::Gzip
        } else if bytes.len() >= 4 && &bytes[..4] == b"\x28\xb5\x2f\xfd" {
            Self::Zstd
        } else if bytes.len() >= 3 && &bytes[..3] == b"BZh" {
            Self::Bzip2
        } else {
            Self::None
        }
    }
}

/// Open a file with transparent decompression. Returns `Send`-safe reader.
pub fn open_maybe_compressed(path: &Path) -> Result<Box<dyn BufRead + Send>> {
    let file = File::open(path).map_err(|e| BunbunError::io(path, e))?;

    let codec = {
        let ext_codec = Codec::from_extension(path);
        if ext_codec != Codec::None {
            ext_codec
        } else {
            let mut peek = [0u8; 4];
            let mut f2 = File::open(path).map_err(|e| BunbunError::io(path, e))?;
            let _ = f2.read(&mut peek);
            Codec::from_magic(&peek)
        }
    };

    let reader: Box<dyn BufRead + Send> = match codec {
        Codec::None => Box::new(BufReader::new(file)),
        Codec::Gzip | Codec::Bgzf => {
            let decoder = GzDecoder::new(file);
            Box::new(BufReader::new(decoder))
        }
        Codec::Zlib => {
            let decoder = ZlibDecoder::new(file);
            Box::new(BufReader::new(decoder))
        }
        Codec::Zstd => {
            let decoder = zstd::Decoder::new(file).map_err(|_| BunbunError::Decompression {
                codec: "zstd",
                detail: "failed to initialize decoder".into(),
            })?;
            Box::new(BufReader::new(decoder))
        }
        Codec::Bzip2 => {
            let decoder = bzip2::read::BzDecoder::new(file);
            Box::new(BufReader::new(decoder))
        }
    };

    Ok(reader)
}

/// Decompress a raw byte slice with a known codec.
pub fn decompress_block(data: &[u8], codec: Codec) -> Result<Vec<u8>> {
    let mut out = Vec::new();
    match codec {
        Codec::None => { out.extend_from_slice(data); }
        Codec::Gzip | Codec::Bgzf => {
            let mut dec = GzDecoder::new(data);
            dec.read_to_end(&mut out).map_err(|_| BunbunError::Decompression {
                codec: "gzip", detail: "block decompression failed".into(),
            })?;
        }
        Codec::Zlib => {
            let mut dec = ZlibDecoder::new(data);
            dec.read_to_end(&mut out).map_err(|_| BunbunError::Decompression {
                codec: "zlib", detail: "block decompression failed".into(),
            })?;
        }
        Codec::Zstd => {
            out = zstd::decode_all(data).map_err(|_| BunbunError::Decompression {
                codec: "zstd", detail: "block decompression failed".into(),
            })?;
        }
        Codec::Bzip2 => {
            let mut dec = bzip2::read::BzDecoder::new(data);
            dec.read_to_end(&mut out).map_err(|_| BunbunError::Decompression {
                codec: "bzip2", detail: "block decompression failed".into(),
            })?;
        }
    }
    Ok(out)
}
