use byteorder::{LittleEndian, ReadBytesExt};
use memmap2::Mmap;
use ndarray::Array2;
use std::fs::File;
use std::io::{BufRead, BufReader, Cursor, Read};
use std::path::{Path, PathBuf};
use std::sync::Arc;

use crate::core::*;
use crate::error::{BunbunError, Result};
use crate::io::{compression, ChunkedReader, Codec, FormatDetector, GenotypeReader};

pub struct BgenReader {
    bgen_path: PathBuf,
    sample_path: Option<PathBuf>,
}

#[derive(Debug, Clone)]
struct BgenHeader {
    offset: u32,
    n_variants: u32,
    n_samples: u32,
    layout: u8,
    compression: BgenCompression,
    has_sample_ids: bool,
}

#[derive(Debug, Clone, Copy)]
enum BgenCompression {
    None,
    Zlib,
    Zstd,
}

impl BgenReader {
    pub fn new(path: impl AsRef<Path>) -> Self {
        let (bgen, sample) = FormatDetector::resolve_bgen(path.as_ref());
        let sample_path = if sample.exists() {
            Some(sample)
        } else {
            None
        };
        Self {
            bgen_path: bgen,
            sample_path,
        }
    }

    pub fn with_sample_file(mut self, path: impl AsRef<Path>) -> Self {
        self.sample_path = Some(path.as_ref().to_path_buf());
        self
    }

    // ── Header ───────────────────────────────────────────────────

    fn parse_header(cursor: &mut Cursor<&[u8]>) -> Result<BgenHeader> {
        let offset = cursor.read_u32::<LittleEndian>().unwrap();
        let header_size = cursor.read_u32::<LittleEndian>().unwrap();
        let n_variants = cursor.read_u32::<LittleEndian>().unwrap();
        let n_samples = cursor.read_u32::<LittleEndian>().unwrap();

        let mut magic = [0u8; 4];
        cursor.read_exact(&mut magic).unwrap();
        if &magic != b"bgen" {
            return Err(BunbunError::BadMagic {
                format: "bgen",
                path: PathBuf::from("<in-memory>"),
            });
        }

        let free_area_size = header_size as usize - 20;
        let mut _skip = vec![0u8; free_area_size];
        cursor.read_exact(&mut _skip).unwrap();

        let flags = cursor.read_u32::<LittleEndian>().unwrap();
        let compression = match flags & 0x3 {
            0 => BgenCompression::None,
            1 => BgenCompression::Zlib,
            2 => BgenCompression::Zstd,
            _ => {
                return Err(BunbunError::UnsupportedVersion {
                    format: "bgen",
                    version: format!("compression flag {}", flags & 0x3),
                })
            }
        };

        let layout = ((flags >> 2) & 0xf) as u8;
        if layout != 1 && layout != 2 {
            return Err(BunbunError::UnsupportedVersion {
                format: "bgen",
                version: format!("layout {layout}"),
            });
        }

        let has_sample_ids = (flags >> 31) & 1 == 1;

        Ok(BgenHeader {
            offset,
            n_variants,
            n_samples,
            layout,
            compression,
            has_sample_ids,
        })
    }

    // ── Sample IDs ───────────────────────────────────────────────

    fn parse_sample_ids(cursor: &mut Cursor<&[u8]>, n_samples: u32) -> Result<Vec<String>> {
        let _block_size = cursor.read_u32::<LittleEndian>().unwrap();
        let n = cursor.read_u32::<LittleEndian>().unwrap();
        assert_eq!(n, n_samples);
        let mut ids = Vec::with_capacity(n as usize);
        for _ in 0..n {
            let len = cursor.read_u16::<LittleEndian>().unwrap() as usize;
            let mut buf = vec![0u8; len];
            cursor.read_exact(&mut buf).unwrap();
            ids.push(String::from_utf8_lossy(&buf).to_string());
        }
        Ok(ids)
    }

    fn parse_sample_file(path: &Path) -> Result<Vec<String>> {
        let file = File::open(path).map_err(|e| BunbunError::io(path, e))?;
        let reader = BufReader::new(file);
        let mut ids = Vec::new();
        let mut line_num = 0;
        for line_result in reader.lines() {
            let line = line_result.map_err(|e| BunbunError::io(path, e))?;
            line_num += 1;
            if line_num <= 2 {
                continue;
            }
            let fields: Vec<&str> = line.split_whitespace().collect();
            if fields.len() >= 2 {
                ids.push(fields[1].to_owned());
            } else if !fields.is_empty() {
                ids.push(fields[0].to_owned());
            }
        }
        Ok(ids)
    }

    // ── Single variant block (layout 2) ──────────────────────────

    fn parse_variant_block_layout2(
        cursor: &mut Cursor<&[u8]>,
        n_samples: usize,
        comp: BgenCompression,
    ) -> Result<(String, String, u32, Allele, Vec<Allele>, Vec<f32>)> {
        let id_len = cursor.read_u16::<LittleEndian>().unwrap() as usize;
        let mut id_buf = vec![0u8; id_len];
        cursor.read_exact(&mut id_buf).unwrap();
        let variant_id = String::from_utf8_lossy(&id_buf).to_string();

        let rsid_len = cursor.read_u16::<LittleEndian>().unwrap() as usize;
        let mut rsid_buf = vec![0u8; rsid_len];
        cursor.read_exact(&mut rsid_buf).unwrap();
        let rsid = String::from_utf8_lossy(&rsid_buf).to_string();

        let chrom_len = cursor.read_u16::<LittleEndian>().unwrap() as usize;
        let mut chrom_buf = vec![0u8; chrom_len];
        cursor.read_exact(&mut chrom_buf).unwrap();
        let chrom = String::from_utf8_lossy(&chrom_buf).to_string();

        let pos = cursor.read_u32::<LittleEndian>().unwrap();
        let n_alleles = cursor.read_u16::<LittleEndian>().unwrap() as usize;

        let mut alleles = Vec::with_capacity(n_alleles);
        for _ in 0..n_alleles {
            let a_len = cursor.read_u32::<LittleEndian>().unwrap() as usize;
            let mut a_buf = vec![0u8; a_len];
            cursor.read_exact(&mut a_buf).unwrap();
            alleles.push(Allele::from(String::from_utf8_lossy(&a_buf).to_string()));
        }

        let ref_allele = alleles.remove(0);
        let alt_alleles = alleles;

        let total_len = cursor.read_u32::<LittleEndian>().unwrap() as usize;

        let dosages = match comp {
            BgenCompression::None => {
                let mut prob_data = vec![0u8; total_len];
                cursor.read_exact(&mut prob_data).unwrap();
                Self::decode_probability_block(&prob_data, n_samples)?
            }
            BgenCompression::Zlib | BgenCompression::Zstd => {
                let _decompressed_len = cursor.read_u32::<LittleEndian>().unwrap() as usize;
                let compressed_len = total_len - 4;
                let mut compressed = vec![0u8; compressed_len];
                cursor.read_exact(&mut compressed).unwrap();
                let codec = match comp {
                    BgenCompression::Zlib => Codec::Zlib,
                    BgenCompression::Zstd => Codec::Zstd,
                    _ => unreachable!(),
                };
                let decompressed = compression::decompress_block(&compressed, codec)?;
                Self::decode_probability_block(&decompressed, n_samples)?
            }
        };

        let id = if rsid == "." || rsid.is_empty() {
            variant_id
        } else {
            rsid
        };
        Ok((chrom, id, pos, ref_allele, alt_alleles, dosages))
    }

    fn decode_probability_block(data: &[u8], n_samples: usize) -> Result<Vec<f32>> {
        if data.len() < 10 + n_samples {
            return Ok(vec![f32::NAN; n_samples]);
        }

        let mut cursor = Cursor::new(data);
        let _n = cursor.read_u32::<LittleEndian>().unwrap() as usize;
        let n_alleles = cursor.read_u16::<LittleEndian>().unwrap();
        let _min_ploidy = cursor.read_u8().unwrap();
        let _max_ploidy = cursor.read_u8().unwrap();

        let mut missingness = vec![false; n_samples];
        for i in 0..n_samples {
            let byte = cursor.read_u8().unwrap();
            missingness[i] = (byte & 0x80) != 0;
        }

        let _phased = cursor.read_u8().unwrap();
        let bits = cursor.read_u8().unwrap() as usize;

        if bits == 0 || n_alleles < 2 {
            return Ok(vec![f32::NAN; n_samples]);
        }

        let max_val = (1u32 << bits) - 1;
        let probs_per_sample = if n_alleles == 2 { 2 } else { n_alleles as usize };
        let remaining = &data[cursor.position() as usize..];
        let mut dosages = Vec::with_capacity(n_samples);
        let mut bit_offset = 0usize;

        for i in 0..n_samples {
            if missingness[i] {
                bit_offset += bits * (probs_per_sample - 1);
                dosages.push(f32::NAN);
                continue;
            }

            let mut probs = Vec::with_capacity(probs_per_sample);
            for _ in 0..(probs_per_sample - 1) {
                let byte_idx = bit_offset / 8;
                let bit_idx = bit_offset % 8;

                if byte_idx >= remaining.len() {
                    probs.push(0.0f32);
                    bit_offset += bits;
                    continue;
                }

                let mut val = 0u32;
                for b in 0..bits {
                    let bi = byte_idx + (bit_idx + b) / 8;
                    let bo = (bit_idx + b) % 8;
                    if bi < remaining.len() && (remaining[bi] >> bo) & 1 == 1 {
                        val |= 1 << b;
                    }
                }
                probs.push(val as f32 / max_val as f32);
                bit_offset += bits;
            }

            if probs.len() >= 2 {
                let p_aa = probs[0];
                let p_ab = probs[1];
                let p_bb = (1.0 - p_aa - p_ab).max(0.0);
                dosages.push(p_ab + 2.0 * p_bb);
            } else if probs.len() == 1 {
                dosages.push(probs[0] * 2.0);
            } else {
                dosages.push(f32::NAN);
            }
        }

        Ok(dosages)
    }

    // ── Open + prepare ───────────────────────────────────────────

    fn open_and_prepare(&self) -> Result<(Mmap, BgenHeader, SampleInfo, u64)> {
        let file =
            File::open(&self.bgen_path).map_err(|e| BunbunError::io(&self.bgen_path, e))?;
        let mmap =
            unsafe { Mmap::map(&file) }.map_err(|e| BunbunError::io(&self.bgen_path, e))?;

        // Use &mmap[..] to get &[u8] for the Cursor
        let mut cursor = Cursor::new(&mmap[..]);
        let header = Self::parse_header(&mut cursor)?;

        let sample_ids = if header.has_sample_ids {
            Self::parse_sample_ids(&mut cursor, header.n_samples)?
        } else if let Some(ref sp) = self.sample_path {
            Self::parse_sample_file(sp)?
        } else {
            (0..header.n_samples as usize)
                .map(|i| format!("sample_{i}"))
                .collect()
        };

        let samples = SampleInfo::from_ids(sample_ids)?;
        let data_start = (header.offset + 4) as u64;

        Ok((mmap, header, samples, data_start))
    }

    fn parse_all(&self) -> Result<(SampleInfo, VariantInfo, DosageMatrix)> {
        let (mmap, header, samples, data_start) = self.open_and_prepare()?;
        let n_samples = header.n_samples as usize;
        let n_variants = header.n_variants as usize;

        let mut cursor = Cursor::new(&mmap[..]);
        cursor.set_position(data_start);

        let mut chroms = Vec::with_capacity(n_variants);
        let mut ids = Vec::with_capacity(n_variants);
        let mut positions = Vec::with_capacity(n_variants);
        let mut ref_alleles = Vec::with_capacity(n_variants);
        let mut alt_alleles_vec = Vec::with_capacity(n_variants);
        let mut matrix = Array2::<f32>::zeros((n_samples, n_variants));

        for var_idx in 0..n_variants {
            if header.layout != 2 {
                return Err(BunbunError::UnsupportedVersion {
                    format: "bgen",
                    version: "layout 1 — use layout 2 (BGEN v1.2+)".into(),
                });
            }
            let (chrom, id, pos, ref_a, alt_a, dosages) =
                Self::parse_variant_block_layout2(&mut cursor, n_samples, header.compression)?;
            chroms.push(chrom);
            ids.push(id);
            positions.push(pos);
            ref_alleles.push(ref_a);
            alt_alleles_vec.push(alt_a);
            for (s, &d) in dosages.iter().enumerate() {
                if s < n_samples {
                    matrix[[s, var_idx]] = d;
                }
            }
        }

        let n = n_variants;
        let variants = VariantInfo::from_parts(
            chroms,
            positions,
            ids,
            ref_alleles,
            alt_alleles_vec,
            vec![f64::NAN; n],
            vec![f64::NAN; n],
            vec!["PASS".to_owned(); n],
        )?;

        Ok((samples, variants, DosageMatrix::new(matrix)))
    }
}

/// Chunked BGEN iterator
struct BgenChunkedIter {
    mmap: Arc<Mmap>,
    header: BgenHeader,
    cursor_pos: u64,
    n_samples: usize,
    n_variants: usize,
    chunk_size: usize,
    variants_read: usize,
}

impl Iterator for BgenChunkedIter {
    type Item = Result<GenotypeChunk>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.variants_read >= self.n_variants {
            return None;
        }

        let start = self.variants_read;
        let end = (start + self.chunk_size).min(self.n_variants);
        let chunk_n = end - start;

        // Use &self.mmap[..] to get &[u8] for the Cursor
        let mut cursor = Cursor::new(&self.mmap[..]);
        cursor.set_position(self.cursor_pos);

        let mut chroms = Vec::with_capacity(chunk_n);
        let mut ids = Vec::with_capacity(chunk_n);
        let mut positions = Vec::with_capacity(chunk_n);
        let mut ref_alleles = Vec::with_capacity(chunk_n);
        let mut alt_alleles = Vec::with_capacity(chunk_n);
        let mut matrix = Array2::<f32>::zeros((self.n_samples, chunk_n));

        for local_idx in 0..chunk_n {
            if self.header.layout != 2 {
                return Some(Err(BunbunError::UnsupportedVersion {
                    format: "bgen",
                    version: "layout 1".into(),
                }));
            }

            let result = BgenReader::parse_variant_block_layout2(
                &mut cursor,
                self.n_samples,
                self.header.compression,
            );

            match result {
                Ok((chrom, id, pos, ref_a, alt_a, dosages)) => {
                    chroms.push(chrom);
                    ids.push(id);
                    positions.push(pos);
                    ref_alleles.push(ref_a);
                    alt_alleles.push(alt_a);
                    for (s, &d) in dosages.iter().enumerate() {
                        if s < self.n_samples {
                            matrix[[s, local_idx]] = d;
                        }
                    }
                }
                Err(e) => return Some(Err(e)),
            }
        }

        self.cursor_pos = cursor.position();
        self.variants_read = end;

        let variants = match VariantInfo::from_parts(
            chroms,
            positions,
            ids,
            ref_alleles,
            alt_alleles,
            vec![f64::NAN; chunk_n],
            vec![f64::NAN; chunk_n],
            vec!["PASS".to_owned(); chunk_n],
        ) {
            Ok(v) => v,
            Err(e) => return Some(Err(e)),
        };

        Some(Ok(GenotypeChunk {
            variants,
            dosages: DosageMatrix::new(matrix),
            offset: start,
            total_variants: self.n_variants,
        }))
    }
}

unsafe impl Send for BgenChunkedIter {}

impl GenotypeReader for BgenReader {
    fn format_name(&self) -> &'static str {
        "bgen"
    }

    fn read_samples(&self) -> Result<SampleInfo> {
        let (_, _, samples, _) = self.open_and_prepare()?;
        Ok(samples)
    }

    fn n_variants(&self) -> Result<usize> {
        let file =
            File::open(&self.bgen_path).map_err(|e| BunbunError::io(&self.bgen_path, e))?;
        let mmap =
            unsafe { Mmap::map(&file) }.map_err(|e| BunbunError::io(&self.bgen_path, e))?;
        let mut cursor = Cursor::new(&mmap[..]);
        let header = Self::parse_header(&mut cursor)?;
        Ok(header.n_variants as usize)
    }

    fn read_all(&self) -> Result<GenotypeData> {
        let (samples, variants, dosages) = self.parse_all()?;
        Ok(GenotypeData {
            samples,
            variants,
            dosages,
            ploidy: Ploidy::Diploid,
            source_format: "bgen",
        })
    }

    fn read_chunked(&self, chunk_size: usize) -> Result<ChunkedReader> {
        let (mmap, header, samples, data_start) = self.open_and_prepare()?;
        let n_samples = header.n_samples as usize;
        let n_variants = header.n_variants as usize;

        let iter = BgenChunkedIter {
            mmap: Arc::new(mmap),
            header,
            cursor_pos: data_start,
            n_samples,
            n_variants,
            chunk_size,
            variants_read: 0,
        };

        Ok(ChunkedReader::new(Box::new(iter), samples, "bgen"))
    }
}
