use flate2::read::{GzDecoder, MultiGzDecoder, ZlibDecoder};
use rayon::prelude::*;
use std::fs::File;
use std::io::{BufRead, BufReader, ErrorKind, 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
        }
    }
}

// ── BGZF parallel decompression ─────────────────────────────────

/// Number of bgzf blocks to read and decompress per batch.
/// 256 blocks × ~64 KB decompressed ≈ 16 MB per batch.
const BGZF_BATCH: usize = 256;

/// Quick check: does this gzip file use the bgzf blocked format?
///
/// Checks for the BC extra-field subfield identifier in the first
/// gzip member header.  Works for all well-formed bgzf files (samtools,
/// htslib, bcftools, etc.).
pub fn is_bgzf(path: &Path) -> bool {
    let Ok(mut f) = File::open(path) else {
        return false;
    };
    let mut header = [0u8; 18];
    if f.read_exact(&mut header).is_err() {
        return false;
    }
    // gzip magic + deflate + FEXTRA flag
    header[0] == 0x1f
        && header[1] == 0x8b
        && header[2] == 0x08
        && (header[3] & 0x04) != 0
        // First extra subfield is BC (the bgzf block-size field)
        && header[12] == 0x42
        && header[13] == 0x43
}

/// Search the gzip extra-field data for the BC subfield and return BSIZE.
///
/// The extra field is a sequence of (SI1, SI2, SLEN, data…) subfields.
/// BGZF's BC subfield has SI1=0x42, SI2=0x43, SLEN=2, and its 2 bytes
/// of data encode BSIZE (total block size minus 1) as little-endian u16.
fn find_bc_bsize(extra: &[u8]) -> std::io::Result<usize> {
    let mut pos = 0;
    while pos + 4 <= extra.len() {
        let si1 = extra[pos];
        let si2 = extra[pos + 1];
        let slen = u16::from_le_bytes([extra[pos + 2], extra[pos + 3]]) as usize;
        if si1 == 0x42 && si2 == 0x43 && slen == 2 && pos + 6 <= extra.len() {
            let bsize = u16::from_le_bytes([extra[pos + 4], extra[pos + 5]]) as usize;
            return Ok(bsize);
        }
        pos += 4 + slen;
    }
    Err(std::io::Error::new(
        ErrorKind::InvalidData,
        "BC subfield not found in bgzf extra field",
    ))
}

/// Read exactly one raw bgzf block (a complete gzip member whose extra
/// field carries the BC/BSIZE tag).  Returns `Ok(None)` at EOF or on
/// encountering the standard 28-byte bgzf EOF marker.
fn read_bgzf_block(file: &mut File) -> std::io::Result<Option<Vec<u8>>> {
    // ── fixed 10-byte gzip header ────────────────────────────────
    let mut fixed = [0u8; 10];
    match file.read_exact(&mut fixed) {
        Ok(()) => {}
        Err(ref e) if e.kind() == ErrorKind::UnexpectedEof => return Ok(None),
        Err(e) => return Err(e),
    }

    if fixed[0] != 0x1f || fixed[1] != 0x8b {
        return Err(std::io::Error::new(
            ErrorKind::InvalidData,
            "invalid bgzf block: bad gzip magic",
        ));
    }
    if (fixed[3] & 0x04) == 0 {
        return Err(std::io::Error::new(
            ErrorKind::InvalidData,
            "invalid bgzf block: FEXTRA flag not set",
        ));
    }

    // ── XLEN + extra field ───────────────────────────────────────
    let mut xlen_buf = [0u8; 2];
    file.read_exact(&mut xlen_buf)?;
    let xlen = u16::from_le_bytes(xlen_buf) as usize;

    let mut extra = vec![0u8; xlen];
    file.read_exact(&mut extra)?;

    let bsize = find_bc_bsize(&extra)?;
    let total_block_size = bsize + 1;
    let header_size = 12 + xlen; // 10 fixed + 2 XLEN + extra

    if total_block_size < header_size {
        return Err(std::io::Error::new(
            ErrorKind::InvalidData,
            "invalid bgzf block: BSIZE smaller than header",
        ));
    }

    // ── assemble full raw block ──────────────────────────────────
    let mut block = vec![0u8; total_block_size];
    block[..10].copy_from_slice(&fixed);
    block[10..12].copy_from_slice(&xlen_buf);
    block[12..12 + xlen].copy_from_slice(&extra);
    let remaining = total_block_size - header_size;
    if remaining > 0 {
        file.read_exact(&mut block[header_size..])?;
    }

    // ── EOF marker check ─────────────────────────────────────────
    // The bgzf EOF marker is an empty gzip member: ISIZE == 0.
    if total_block_size >= 4 {
        let off = total_block_size - 4;
        let isize_val = u32::from_le_bytes([block[off], block[off + 1], block[off + 2], block[off + 3]]);
        if isize_val == 0 && remaining <= 16 {
            // Empty block – treat as EOF
            return Ok(None);
        }
    }

    Ok(Some(block))
}

/// A Read adapter that reads raw bgzf blocks in batches and decompresses
/// them in parallel using rayon.
///
/// Each bgzf block is an independent gzip member (≤ 64 KiB decompressed),
/// so they can be inflated concurrently.  The reader maintains an internal
/// buffer of decompressed bytes and refills it when exhausted.
pub struct BgzfParallelReader {
    file: File,
    buf: Vec<u8>,
    pos: usize,
    eof: bool,
}

impl BgzfParallelReader {
    pub fn new(file: File) -> Self {
        Self {
            file,
            buf: Vec::new(),
            pos: 0,
            eof: false,
        }
    }

    /// Read up to BGZF_BATCH raw blocks and decompress them in parallel.
    fn refill(&mut self) -> std::io::Result<()> {
        // Collect raw blocks
        let mut raw_blocks = Vec::with_capacity(BGZF_BATCH);
        for _ in 0..BGZF_BATCH {
            match read_bgzf_block(&mut self.file)? {
                Some(block) => raw_blocks.push(block),
                None => {
                    self.eof = true;
                    break;
                }
            }
        }

        if raw_blocks.is_empty() {
            self.buf.clear();
            self.pos = 0;
            return Ok(());
        }

        // Decompress in parallel – each block is an independent gzip member
        let decompressed: Vec<std::io::Result<Vec<u8>>> = raw_blocks
            .par_iter()
            .map(|block| {
                let mut dec = GzDecoder::new(&block[..]);
                let mut out = Vec::with_capacity(65536);
                dec.read_to_end(&mut out)?;
                Ok(out)
            })
            .collect();

        // Concatenate into a single contiguous buffer (preserves block order)
        let total: usize = decompressed
            .iter()
            .filter_map(|r| r.as_ref().ok())
            .map(|v| v.len())
            .sum();

        self.buf.clear();
        self.buf.reserve(total);
        for result in decompressed {
            self.buf.extend_from_slice(&result?);
        }
        self.pos = 0;

        Ok(())
    }
}

impl Read for BgzfParallelReader {
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        loop {
            let avail = self.buf.len() - self.pos;
            if avail > 0 {
                let n = buf.len().min(avail);
                buf[..n].copy_from_slice(&self.buf[self.pos..self.pos + n]);
                self.pos += n;
                return Ok(n);
            }
            if self.eof {
                return Ok(0);
            }
            self.refill()?;
            if self.buf.is_empty() {
                return Ok(0);
            }
        }
    }
}

// ── Public API ───────────────────────────────────────────────────

/// Open a file with transparent decompression. Returns `Send`-safe reader.
///
/// For bgzf files (the standard for .vcf.gz in bioinformatics) this uses
/// parallel block decompression via rayon.  Plain gzip, zstd, bzip2 fall
/// back to sequential decoders.
pub fn open_maybe_compressed(path: &Path) -> Result<Box<dyn BufRead + Send>> {
    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 => {
            let file = File::open(path).map_err(|e| BunbunError::io(path, e))?;
            Box::new(BufReader::new(file))
        }
        Codec::Gzip | Codec::Bgzf => {
            if is_bgzf(path) {
                // Parallel bgzf decompression – the fast path
                let file = File::open(path).map_err(|e| BunbunError::io(path, e))?;
                let par_reader = BgzfParallelReader::new(file);
                // 256 KB BufReader – the internal buffer is already large,
                // but this smooths out line-boundary reads.
                Box::new(BufReader::with_capacity(262_144, par_reader))
            } else {
                // Plain (non-blocked) gzip – sequential fallback
                let file = File::open(path).map_err(|e| BunbunError::io(path, e))?;
                let decoder = MultiGzDecoder::new(file);
                Box::new(BufReader::new(decoder))
            }
        }
        Codec::Zlib => {
            let file = File::open(path).map_err(|e| BunbunError::io(path, e))?;
            let decoder = ZlibDecoder::new(file);
            Box::new(BufReader::new(decoder))
        }
        Codec::Zstd => {
            let file = File::open(path).map_err(|e| BunbunError::io(path, e))?;
            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 file = File::open(path).map_err(|e| BunbunError::io(path, e))?;
            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 = MultiGzDecoder::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)
}
