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

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

pub struct PlinkReader {
    bed: PathBuf,
    bim: PathBuf,
    fam: PathBuf,
}

impl PlinkReader {
    pub fn new(path: impl AsRef<Path>) -> Self {
        let (bed, bim, fam) = FormatDetector::resolve_plink(path.as_ref());
        Self { bed, bim, fam }
    }

    pub fn from_paths(bed: PathBuf, bim: PathBuf, fam: PathBuf) -> Self {
        Self { bed, bim, fam }
    }

    // ── .fam ─────────────────────────────────────────────────────

    fn parse_fam(&self) -> Result<SampleInfo> {
        let file = File::open(&self.fam).map_err(|e| BunbunError::io(&self.fam, e))?;
        let reader = BufReader::new(file);

        let mut fid = Vec::new();
        let mut iid = Vec::new();
        let mut father = Vec::new();
        let mut mother = Vec::new();
        let mut sex = Vec::new();
        let mut pheno = Vec::new();

        for (i, line) in reader.lines().enumerate() {
            let line = line.map_err(|e| BunbunError::io(&self.fam, e))?;
            let line = line.trim();
            if line.is_empty() {
                continue;
            }
            let fields: Vec<&str> = line.split_whitespace().collect();
            if fields.len() < 6 {
                return Err(BunbunError::ParseRecord {
                    format: "plink",
                    record: i,
                    detail: format!(".fam line has {} fields, expected 6", fields.len()),
                });
            }
            fid.push(fields[0].to_owned());
            iid.push(fields[1].to_owned());
            father.push(fields[2].to_owned());
            mother.push(fields[3].to_owned());
            sex.push(Sex::from_plink(fields[4].parse::<i8>().unwrap_or(0)));
            pheno.push(fields[5].parse::<f64>().unwrap_or(f64::NAN));
        }

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

    // ── .bim ─────────────────────────────────────────────────────

    fn parse_bim(&self) -> Result<VariantInfo> {
        let (chrom, id, cm, pos, a1, a2) = self.parse_bim_raw()?;
        VariantInfo::from_bim(chrom, id, cm, pos, a1, a2)
    }

    fn parse_bim_raw(
        &self,
    ) -> Result<(Vec<String>, Vec<String>, Vec<f64>, Vec<u32>, Vec<String>, Vec<String>)> {
        let file = File::open(&self.bim).map_err(|e| BunbunError::io(&self.bim, e))?;
        let reader = BufReader::new(file);

        let mut chrom = Vec::new();
        let mut id = Vec::new();
        let mut cm = Vec::new();
        let mut pos = Vec::new();
        let mut a1 = Vec::new();
        let mut a2 = Vec::new();

        for (i, line) in reader.lines().enumerate() {
            let line = line.map_err(|e| BunbunError::io(&self.bim, e))?;
            let line = line.trim();
            if line.is_empty() {
                continue;
            }
            let fields: Vec<&str> = line.split_whitespace().collect();
            if fields.len() < 6 {
                return Err(BunbunError::ParseRecord {
                    format: "plink",
                    record: i,
                    detail: format!(".bim line has {} fields, expected 6", fields.len()),
                });
            }
            chrom.push(fields[0].to_owned());
            id.push(fields[1].to_owned());
            cm.push(fields[2].parse::<f64>().unwrap_or(0.0));
            pos.push(fields[3].parse::<u32>().unwrap_or(0));
            a1.push(fields[4].to_owned());
            a2.push(fields[5].to_owned());
        }

        Ok((chrom, id, cm, pos, a1, a2))
    }

    // ── .bed ─────────────────────────────────────────────────────

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

        if mmap.len() < 3 || mmap[0] != 0x6c || mmap[1] != 0x1b {
            return Err(BunbunError::BadMagic {
                format: "plink",
                path: self.bed.clone(),
            });
        }
        if mmap[2] != 0x01 {
            return Err(BunbunError::UnsupportedVersion {
                format: "plink",
                version: "sample-major .bed (mode byte != 0x01)".into(),
            });
        }

        let bytes_per_variant = (n_samples + 3) / 4;
        let expected_size = 3 + bytes_per_variant * n_variants;
        if mmap.len() < expected_size {
            return Err(BunbunError::DimensionMismatch {
                expected: format!("{expected_size} bytes"),
                got: format!("{} bytes", mmap.len()),
                context: ".bed file too small for n_samples × n_variants",
            });
        }

        Ok(mmap)
    }

    fn parse_bed(&self, n_samples: usize, n_variants: usize) -> Result<DosageMatrix> {
        let mmap = self.open_bed(n_samples, n_variants)?;
        decode_bed_range(&mmap, n_samples, n_variants, 0, n_variants)
    }

    fn chunked_iter(&self, chunk_size: usize) -> Result<(SampleInfo, PlinkChunkedIter)> {
        let samples = self.parse_fam()?;
        let (chrom, id, cm, pos, a1, a2) = self.parse_bim_raw()?;
        let n_samples = samples.n_samples();
        let n_variants = chrom.len();
        let mmap = Arc::new(self.open_bed(n_samples, n_variants)?);

        let iter = PlinkChunkedIter {
            mmap,
            chrom,
            id,
            cm,
            pos,
            a1,
            a2,
            n_samples,
            n_variants,
            chunk_size,
            cursor: 0,
        };

        Ok((samples, iter))
    }
}

/// Decode a range [start, end) of variants from a .bed mmap. Uses rayon.
fn decode_bed_range(
    mmap: &Mmap,
    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;
    let bed_data = &mmap[3..];
    let mut matrix = Array2::<f32>::zeros((n_samples, chunk_n));

    // Index-based parallel iteration — avoids AxisIterMut trait bound issues
    let raw = matrix.as_slice_mut().expect("matrix is contiguous");
    (0..chunk_n).into_par_iter().for_each(|local_idx| {
        let global_idx = start + local_idx;
        let offset = global_idx * bytes_per_variant;
        let variant_bytes = &bed_data[offset..offset + bytes_per_variant];
        // Each column in a row-major (n_samples × chunk_n) matrix:
        // element [i, j] is at index i * chunk_n + j
        // SAFETY: each parallel task writes to non-overlapping column positions
        for s in 0..n_samples {
            let byte_idx = s / 4;
            let bit_offset = (s % 4) * 2;
            let code = (variant_bytes[byte_idx] >> bit_offset) & 0b11;
            let dosage = GenotypeValue::from_plink_code(code).to_dosage();
            unsafe {
                let ptr = raw.as_ptr() as *mut f32;
                *ptr.add(s * chunk_n + local_idx) = dosage;
            }
        }
    });

    Ok(DosageMatrix::new(matrix))
}

struct PlinkChunkedIter {
    mmap: Arc<Mmap>,
    chrom: Vec<String>,
    id: Vec<String>,
    cm: Vec<f64>,
    pos: Vec<u32>,
    a1: Vec<String>,
    a2: Vec<String>,
    n_samples: usize,
    n_variants: usize,
    chunk_size: usize,
    cursor: usize,
}

impl Iterator for PlinkChunkedIter {
    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);
        self.cursor = end;

        let variants = match VariantInfo::from_bim(
            self.chrom[start..end].to_vec(),
            self.id[start..end].to_vec(),
            self.cm[start..end].to_vec(),
            self.pos[start..end].to_vec(),
            self.a1[start..end].to_vec(),
            self.a2[start..end].to_vec(),
        ) {
            Ok(v) => v,
            Err(e) => return Some(Err(e)),
        };

        let dosages = match decode_bed_range(
            &self.mmap,
            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 PlinkChunkedIter {}

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

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

    fn n_variants(&self) -> Result<usize> {
        let file = File::open(&self.bim).map_err(|e| BunbunError::io(&self.bim, e))?;
        let reader = BufReader::new(file);
        Ok(reader
            .lines()
            .filter(|l| l.as_ref().map(|s| !s.trim().is_empty()).unwrap_or(false))
            .count())
    }

    fn read_all(&self) -> Result<GenotypeData> {
        let samples = self.parse_fam()?;
        let variants = self.parse_bim()?;
        let dosages = self.parse_bed(samples.n_samples(), variants.n_variants())?;

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

    fn read_chunked(&self, chunk_size: usize) -> Result<ChunkedReader> {
        let (samples, iter) = self.chunked_iter(chunk_size)?;
        Ok(ChunkedReader::new(Box::new(iter), samples, "plink"))
    }
}
