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

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

pub struct Plink2Reader {
    pgen: PathBuf,
    pvar: PathBuf,
    psam: PathBuf,
}

impl Plink2Reader {
    pub fn new(path: impl AsRef<Path>) -> Self {
        let (pgen, pvar, psam) = FormatDetector::resolve_plink2(path.as_ref());
        Self { pgen, pvar, psam }
    }

    pub fn from_paths(pgen: PathBuf, pvar: PathBuf, psam: PathBuf) -> Self {
        Self { pgen, pvar, psam }
    }

    // ── .psam ────────────────────────────────────────────────────

    fn parse_psam(&self) -> Result<SampleInfo> {
        let reader = open_maybe_compressed(&self.psam)?;
        let mut header_cols: Vec<String> = Vec::new();
        let mut rows: Vec<Vec<String>> = Vec::new();

        for line_result in reader.lines() {
            let line = line_result.map_err(|e| BunbunError::io(&self.psam, e))?;
            let line = line.trim().to_string();
            if line.is_empty() {
                continue;
            }
            if line.starts_with("#FID")
                || line.starts_with("#IID")
                || line.starts_with("FID")
                || line.starts_with("IID")
            {
                let clean = line.trim_start_matches('#');
                header_cols = clean.split('\t').map(|s| s.to_owned()).collect();
                continue;
            }
            if line.starts_with('#') {
                continue;
            }
            let fields: Vec<String> = line.split('\t').map(|s| s.to_owned()).collect();
            rows.push(fields);
        }

        let col_idx =
            |name: &str| -> Option<usize> { header_cols.iter().position(|c| c.eq_ignore_ascii_case(name)) };

        let iid_col = col_idx("IID").unwrap_or(0);
        let fid_col = col_idx("FID");
        let sex_col = col_idx("SEX");
        let pat_col = col_idx("PAT");
        let mat_col = col_idx("MAT");
        let pheno_col = col_idx("PHENO1");

        let get = |row: &[String], idx: Option<usize>, default: &str| -> String {
            idx.and_then(|i| row.get(i))
                .map(|s| s.to_owned())
                .unwrap_or_else(|| default.to_owned())
        };

        let fid: Vec<String> = rows.iter().map(|r| get(r, fid_col, ".")).collect();
        let iid: Vec<String> = rows.iter().map(|r| r.get(iid_col).cloned().unwrap_or_default()).collect();
        let father: Vec<String> = rows.iter().map(|r| get(r, pat_col, "0")).collect();
        let mother: Vec<String> = rows.iter().map(|r| get(r, mat_col, "0")).collect();
        let sex: Vec<Sex> = rows
            .iter()
            .map(|r| {
                sex_col
                    .and_then(|i| r.get(i))
                    .map(|s| Sex::from_plink(s.parse::<i8>().unwrap_or(0)))
                    .unwrap_or(Sex::Unknown)
            })
            .collect();
        let pheno: Vec<f64> = rows
            .iter()
            .map(|r| {
                pheno_col
                    .and_then(|i| r.get(i))
                    .and_then(|s| s.parse::<f64>().ok())
                    .unwrap_or(f64::NAN)
            })
            .collect();

        SampleInfo::from_plink_fam(fid, iid, father, mother, sex, pheno)
    }

    // ── .pvar ────────────────────────────────────────────────────

    fn parse_pvar(&self) -> Result<VariantInfo> {
        let raw = self.parse_pvar_raw()?;
        let n = raw.chrom.len();
        VariantInfo::from_parts(
            raw.chrom,
            raw.pos,
            raw.id,
            raw.ref_allele,
            raw.alt_alleles,
            vec![f64::NAN; n],
            vec![f64::NAN; n],
            vec!["PASS".to_owned(); n],
        )
    }

    fn parse_pvar_raw(&self) -> Result<PvarRaw> {
        let reader = open_maybe_compressed(&self.pvar)?;
        let mut header_cols: Vec<String> = Vec::new();
        let mut raw = PvarRaw::default();

        for line_result in reader.lines() {
            let line = line_result.map_err(|e| BunbunError::io(&self.pvar, e))?;
            let line = line.trim().to_string();
            if line.is_empty() || (line.starts_with("##") && !line.starts_with("#CHROM")) {
                continue;
            }
            if line.starts_with("#CHROM") || line.starts_with("CHROM") {
                let clean = line.trim_start_matches('#');
                header_cols = clean.split('\t').map(|s| s.to_owned()).collect();
                continue;
            }

            let fields: Vec<&str> = line.split('\t').collect();
            let col_idx =
                |name: &str| -> Option<usize> { header_cols.iter().position(|c| c.eq_ignore_ascii_case(name)) };

            let chrom_col = col_idx("CHROM").unwrap_or(0);
            let pos_col = col_idx("POS").unwrap_or(1);
            let id_col = col_idx("ID").unwrap_or(2);
            let ref_col = col_idx("REF").unwrap_or(3);
            let alt_col = col_idx("ALT").unwrap_or(4);

            raw.chrom.push(fields.get(chrom_col).unwrap_or(&".").to_string());
            raw.pos.push(
                fields
                    .get(pos_col)
                    .and_then(|s| s.parse().ok())
                    .unwrap_or(0),
            );
            let vid = fields.get(id_col).unwrap_or(&".");
            raw.id.push(if *vid == "." {
                format!(
                    "{}:{}",
                    fields.get(chrom_col).unwrap_or(&"."),
                    fields.get(pos_col).unwrap_or(&"0")
                )
            } else {
                vid.to_string()
            });
            raw.ref_allele
                .push(Allele::new(fields.get(ref_col).unwrap_or(&"N")));
            raw.alt_alleles.push(
                fields
                    .get(alt_col)
                    .unwrap_or(&".")
                    .split(',')
                    .map(Allele::new)
                    .collect(),
            );
        }

        Ok(raw)
    }

    // ── .pgen ────────────────────────────────────────────────────

    fn open_pgen(&self, n_samples: usize, n_variants: usize) -> Result<(Mmap, u8)> {
        let file = File::open(&self.pgen).map_err(|e| BunbunError::io(&self.pgen, e))?;
        let mmap = unsafe { Mmap::map(&file) }.map_err(|e| BunbunError::io(&self.pgen, e))?;

        if mmap.len() < 12 || &mmap[0..2] != b"l\x1b" {
            return Err(BunbunError::BadMagic {
                format: "plink2",
                path: self.pgen.clone(),
            });
        }

        let storage_mode = mmap[2];
        let mut cursor = Cursor::new(&mmap[3..11]);
        let pgen_nv = cursor.read_u32::<LittleEndian>().unwrap() as usize;
        let pgen_ns = cursor.read_u32::<LittleEndian>().unwrap() as usize;

        if pgen_ns != n_samples {
            return Err(BunbunError::DimensionMismatch {
                expected: format!("{n_samples} samples (from .psam)"),
                got: format!("{pgen_ns} (in .pgen header)"),
                context: ".pgen/.psam sample count mismatch",
            });
        }
        if pgen_nv != n_variants {
            return Err(BunbunError::DimensionMismatch {
                expected: format!("{n_variants} variants (from .pvar)"),
                got: format!("{pgen_nv} (in .pgen header)"),
                context: ".pgen/.pvar variant count mismatch",
            });
        }

        if storage_mode != 0x02 && storage_mode != 0x03 {
            return Err(BunbunError::UnsupportedVersion {
                format: "plink2",
                version: format!("storage mode 0x{storage_mode:02x}"),
            });
        }

        Ok((mmap, storage_mode))
    }

    fn parse_pgen(&self, n_samples: usize, n_variants: usize) -> Result<DosageMatrix> {
        let (mmap, mode) = self.open_pgen(n_samples, n_variants)?;
        decode_pgen_range(&mmap, mode, n_samples, n_variants, 0, n_variants)
    }
}

#[derive(Default)]
struct PvarRaw {
    chrom: Vec<String>,
    pos: Vec<u32>,
    id: Vec<String>,
    ref_allele: Vec<Allele>,
    alt_alleles: Vec<Vec<Allele>>,
}

/// Decode a range of variants from a .pgen mmap.
fn decode_pgen_range(
    mmap: &Mmap,
    mode: u8,
    n_samples: usize,
    total_variants: usize,
    start: usize,
    end: usize,
) -> Result<DosageMatrix> {
    let chunk_n = end - start;
    let bytes_per_variant = (n_samples + 3) / 4;

    match mode {
        0x02 => {
            let header_size = 12;
            let data = &mmap[header_size..];
            let mut matrix = Array2::<f32>::zeros((n_samples, chunk_n));

            // Index-based parallel iteration
            let raw = matrix.as_slice_mut().expect("contiguous");
            (0..chunk_n).into_par_iter().for_each(|local_idx: usize| {
                let global_idx = start + local_idx;
                let offset = global_idx * bytes_per_variant;
                let chunk = &data[offset..offset + bytes_per_variant];
                for s in 0..n_samples {
                    let byte_idx = s / 4;
                    let bit_offset = (s % 4) * 2;
                    let code = (chunk[byte_idx] >> bit_offset) & 0b11;
                    let dosage = match code {
                        0 => 0.0f32,
                        1 => 1.0,
                        2 => 2.0,
                        _ => f32::NAN,
                    };
                    unsafe {
                        let ptr = raw.as_ptr() as *mut f32;
                        *ptr.add(s * chunk_n + local_idx) = dosage;
                    }
                }
            });

            Ok(DosageMatrix::new(matrix))
        }
        0x03 => {
            let header_size = 12;
            let offset_table_size = total_variants * 8;
            let mut cursor = Cursor::new(&mmap[header_size..header_size + offset_table_size]);
            let mut offsets = Vec::with_capacity(total_variants);
            for _ in 0..total_variants {
                offsets.push(cursor.read_u64::<LittleEndian>().unwrap() as usize);
            }

            let mut matrix = Array2::<f32>::zeros((n_samples, chunk_n));
            let raw = matrix.as_slice_mut().expect("contiguous");

            (0..chunk_n).into_par_iter().for_each(|local_idx: usize| {
                let global_idx = start + local_idx;
                let off = offsets[global_idx];
                for s in 0..n_samples {
                    let dosage = if off + bytes_per_variant > mmap.len() {
                        f32::NAN
                    } else {
                        let chunk = &mmap[off..off + bytes_per_variant];
                        let byte_idx = s / 4;
                        let bit_offset = (s % 4) * 2;
                        let code = (chunk[byte_idx] >> bit_offset) & 0b11;
                        match code {
                            0 => 0.0f32,
                            1 => 1.0,
                            2 => 2.0,
                            _ => f32::NAN,
                        }
                    };
                    unsafe {
                        let ptr = raw.as_ptr() as *mut f32;
                        *ptr.add(s * chunk_n + local_idx) = dosage;
                    }
                }
            });

            Ok(DosageMatrix::new(matrix))
        }
        _ => unreachable!(),
    }
}

struct Plink2ChunkedIter {
    mmap: Arc<Mmap>,
    pvar: PvarRaw,
    mode: u8,
    n_samples: usize,
    n_variants: usize,
    chunk_size: usize,
    cursor: usize,
}

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

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

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

        let variants = match VariantInfo::from_parts(
            self.pvar.chrom[start..end].to_vec(),
            self.pvar.pos[start..end].to_vec(),
            self.pvar.id[start..end].to_vec(),
            self.pvar.ref_allele[start..end].to_vec(),
            self.pvar.alt_alleles[start..end].to_vec(),
            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)),
        };

        let dosages = match decode_pgen_range(
            &self.mmap,
            self.mode,
            self.n_samples,
            self.n_variants,
            start,
            end,
        ) {
            Ok(d) => d,
            Err(e) => return Some(Err(e)),
        };

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

unsafe impl Send for Plink2ChunkedIter {}

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

    fn read_samples(&self) -> Result<SampleInfo> {
        self.parse_psam()
    }

    fn n_variants(&self) -> Result<usize> {
        Ok(self.parse_pvar_raw()?.chrom.len())
    }

    fn read_all(&self) -> Result<GenotypeData> {
        let samples = self.parse_psam()?;
        let variants = self.parse_pvar()?;
        let dosages = self.parse_pgen(samples.n_samples(), variants.n_variants())?;

        Ok(GenotypeData {
            samples,
            variants,
            dosages,
            ploidy: Ploidy::Diploid,
            source_format: "plink2",
        })
    }

    fn read_chunked(&self, chunk_size: usize) -> Result<ChunkedReader> {
        let samples = self.parse_psam()?;
        let pvar = self.parse_pvar_raw()?;
        let n_samples = samples.n_samples();
        let n_variants = pvar.chrom.len();
        let (mmap, mode) = self.open_pgen(n_samples, n_variants)?;

        let iter = Plink2ChunkedIter {
            mmap: Arc::new(mmap),
            pvar,
            mode,
            n_samples,
            n_variants,
            chunk_size,
            cursor: 0,
        };

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