use ndarray::Array2;
use std::io::BufRead;
use std::path::{Path, PathBuf};

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

pub struct VcfReader {
    path: PathBuf,
}

impl VcfReader {
    pub fn new(path: impl AsRef<Path>) -> Self {
        Self {
            path: path.as_ref().to_path_buf(),
        }
    }

    fn parse_header(reader: &mut dyn BufRead) -> Result<Vec<String>> {
        let mut sample_ids: Vec<String> = Vec::new();
        let mut line = String::new();

        loop {
            line.clear();
            let bytes_read = reader.read_line(&mut line).map_err(|e| BunbunError::Io {
                path: PathBuf::from("<vcf>"),
                source: e,
            })?;
            if bytes_read == 0 {
                break;
            }

            if line.starts_with("##") {
                continue;
            }

            if line.starts_with("#CHROM") {
                let fields: Vec<&str> = line.trim_end().split('\t').collect();
                if fields.len() > 9 {
                    sample_ids = fields[9..].iter().map(|s| s.to_string()).collect();
                }
                break;
            }
        }

        if sample_ids.is_empty() {
            return Err(BunbunError::MalformedHeader {
                format: "vcf",
                offset: 0,
                detail: "no sample columns found in #CHROM header".into(),
            });
        }

        Ok(sample_ids)
    }

    /// Original parse_record – used by read_all().
    /// Kept for backward compatibility; read_all is not the hot path.
    fn parse_record(
        line: &str,
        n_samples: usize,
    ) -> Result<Option<(String, u32, String, Allele, Vec<Allele>, f64, String, Vec<f32>)>> {
        let line = line.trim_end();
        if line.is_empty() {
            return Ok(None);
        }

        let fields: Vec<&str> = line.splitn(10, '\t').collect();
        if fields.len() < 10 {
            return Ok(None);
        }

        let chrom = fields[0].to_owned();
        let pos = fields[1].parse::<u32>().unwrap_or(0);
        let id = if fields[2] == "." {
            format!("{}:{}:{}:{}", fields[0], fields[1], fields[3], fields[4])
        } else {
            fields[2].to_owned()
        };
        let ref_allele = Allele::new(fields[3]);
        let alt_alleles: Vec<Allele> = fields[4].split(',').map(Allele::new).collect();
        let qual = fields[5].parse::<f64>().unwrap_or(f64::NAN);
        let filter = fields[6].to_owned();

        let format_field = fields[8];
        let gt_idx = find_format_index(format_field, "GT");
        let ds_idx = find_format_index(format_field, "DS");

        let sample_block = fields[9];
        let mut dosages = vec![f32::NAN; n_samples];

        // Use the fast byte-level parser even in the read_all path
        parse_samples_into(
            sample_block.as_bytes(),
            n_samples,
            gt_idx,
            ds_idx,
            &mut dosages,
        );

        Ok(Some((
            chrom, pos, id, ref_allele, alt_alleles, qual, filter, dosages,
        )))
    }

    fn count_variants(path: &Path) -> Result<usize> {
        let reader = open_maybe_compressed(path)?;
        let mut count = 0;
        for line_result in reader.lines() {
            let line = line_result.map_err(|e| BunbunError::io(path, e))?;
            if !line.starts_with('#') && !line.trim().is_empty() {
                count += 1;
            }
        }
        Ok(count)
    }

    fn parse_all(&self) -> Result<(SampleInfo, VariantInfo, DosageMatrix)> {
        let mut reader = open_maybe_compressed(&self.path)?;
        let sample_ids = Self::parse_header(&mut *reader)?;
        let n_samples = sample_ids.len();
        let samples = SampleInfo::from_ids(sample_ids)?;

        let mut chroms = Vec::new();
        let mut positions = Vec::new();
        let mut ids = Vec::new();
        let mut refs = Vec::new();
        let mut alts = Vec::new();
        let mut quals = Vec::new();
        let mut filters = Vec::new();
        let mut all_dosages: Vec<Vec<f32>> = Vec::new();

        let mut line = String::new();
        loop {
            line.clear();
            let n = reader
                .read_line(&mut line)
                .map_err(|e| BunbunError::io(&self.path, e))?;
            if n == 0 {
                break;
            }

            if let Some((chrom, pos, id, ref_a, alt_a, qual, filt, dosages)) =
                Self::parse_record(&line, n_samples)?
            {
                chroms.push(chrom);
                positions.push(pos);
                ids.push(id);
                refs.push(ref_a);
                alts.push(alt_a);
                quals.push(qual);
                filters.push(filt);
                all_dosages.push(dosages);
            }
        }

        let n_variants = chroms.len();
        let mut matrix = Array2::<f32>::zeros((n_samples, n_variants));
        for (j, row) in all_dosages.iter().enumerate() {
            for (i, &val) in row.iter().enumerate() {
                matrix[[i, j]] = val;
            }
        }

        let cm = vec![f64::NAN; n_variants];
        let variants =
            VariantInfo::from_parts(chroms, positions, ids, refs, alts, cm, quals, filters)?;

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

// ── Zero-alloc byte-level parsing helpers ────────────────────────

/// Convert a byte slice to a String (lossy but fast for ASCII VCF data).
#[inline]
fn string_from_bytes(b: &[u8]) -> String {
    String::from_utf8_lossy(b).into_owned()
}

/// Fast u32 parser for ASCII digits – avoids str conversion + parse.
#[inline]
fn parse_u32_fast(b: &[u8]) -> u32 {
    let mut n: u32 = 0;
    for &ch in b {
        if ch >= b'0' && ch <= b'9' {
            n = n.wrapping_mul(10).wrapping_add((ch - b'0') as u32);
        } else {
            break;
        }
    }
    n
}

/// Fast f64 parser – delegates to std for correctness.
#[inline]
fn parse_f64_fast(b: &[u8]) -> f64 {
    std::str::from_utf8(b)
        .ok()
        .and_then(|s| s.parse::<f64>().ok())
        .unwrap_or(f64::NAN)
}

#[inline]
fn find_format_index(format: &str, key: &str) -> Option<usize> {
    format.split(':').position(|f| f == key)
}

/// Byte-level FORMAT index search – avoids &str conversion.
#[inline]
fn find_format_index_bytes(format: &[u8], key: &[u8]) -> Option<usize> {
    let mut idx = 0usize;
    let mut start = 0;
    let len = format.len();
    while start <= len {
        let end = memchr::memchr(b':', &format[start..])
            .map(|o| start + o)
            .unwrap_or(len);
        if &format[start..end] == key {
            return Some(idx);
        }
        idx += 1;
        start = end + 1;
    }
    None
}

/// Parse a VCF GT subfield directly from bytes into a dosage float.
///
/// Dosage = count of non-reference alleles (allele index != 0).
///
/// Handles diploid (0/0, 0/1, 1/1, 0/2, 2/2, 1|0, …), haploid (0, 1),
/// and multi-allelic GTs without any allocation.
///
/// Examples:
///   0/0 → 0.0    0/1 → 1.0    1/1 → 2.0
///   0/2 → 1.0    2/2 → 2.0    1/2 → 2.0
///   .   → NAN    ./. → NAN
#[inline]
fn gt_dosage_from_bytes(gt: &[u8]) -> f32 {
    // Diploid fast path (covers >99% of real data):
    // Single-digit allele indices separated by / or |
    if gt.len() >= 3 {
        let a = gt[0];
        let sep = gt[1];
        let b = gt[2];
        if a == b'.' || b == b'.' {
            return f32::NAN;
        }
        if (sep == b'/' || sep == b'|') && a >= b'0' && b >= b'0' {
            // Count non-reference alleles (anything != '0')
            return (a != b'0') as u8 as f32 + (b != b'0') as u8 as f32;
        }
    }
    // Haploid (e.g. chrX male: "0" or "1")
    if gt.len() == 1 {
        if gt[0] == b'.' {
            return f32::NAN;
        }
        if gt[0] >= b'0' && gt[0] <= b'9' {
            return (gt[0] != b'0') as u8 as f32;
        }
    }
    // General path: multi-digit allele indices, polyploid, etc.
    // Count non-reference alleles (index != 0).
    let mut dosage = 0u32;
    let mut current = 0u32;
    let mut has_digit = false;
    for &ch in gt {
        if ch == b'.' {
            return f32::NAN;
        } else if ch >= b'0' && ch <= b'9' {
            current = current * 10 + (ch - b'0') as u32;
            has_digit = true;
        } else if ch == b'/' || ch == b'|' {
            if has_digit && current != 0 {
                dosage += 1;
            }
            current = 0;
            has_digit = false;
        }
    }
    // Last allele after final separator
    if has_digit && current != 0 {
        dosage += 1;
    }
    dosage as f32
}

/// Navigate to the `n`th colon-separated subfield in a byte slice.
#[inline]
fn nth_subfield(data: &[u8], n: usize) -> Option<&[u8]> {
    let mut start = 0;
    for _ in 0..n {
        match memchr::memchr(b':', &data[start..]) {
            Some(offset) => start += offset + 1,
            None => return None,
        }
    }
    let end = memchr::memchr(b':', &data[start..])
        .map(|o| start + o)
        .unwrap_or(data.len());
    Some(&data[start..end])
}

/// Parse all sample genotypes from the sample block directly into a
/// pre-allocated output slice.  Uses `memchr` for tab/colon boundaries.
fn parse_samples_into(
    sample_data: &[u8],
    n_samples: usize,
    gt_idx: Option<usize>,
    ds_idx: Option<usize>,
    out: &mut [f32],
) {
    // Trim trailing newlines/carriage returns
    let mut end = sample_data.len();
    while end > 0 && (sample_data[end - 1] == b'\n' || sample_data[end - 1] == b'\r') {
        end -= 1;
    }
    let data = &sample_data[..end];
    let data_len = data.len();

    let mut sample_i = 0;
    let mut pos = 0;

    // Fast path: GT is the first subfield (GT index == 0, no DS).
    // This is by far the most common layout (FORMAT = GT:…).
    // We avoid the nth_subfield call entirely.
    if ds_idx.is_none() && gt_idx == Some(0) {
        while sample_i < n_samples && pos < data_len {
            // Find end of this sample (next tab or end of data)
            let sample_end = memchr::memchr(b'\t', &data[pos..])
                .map(|o| pos + o)
                .unwrap_or(data_len);

            // GT is before the first colon (or the entire field if no colon)
            let gt_end = memchr::memchr(b':', &data[pos..sample_end])
                .map(|o| pos + o)
                .unwrap_or(sample_end);

            out[sample_i] = gt_dosage_from_bytes(&data[pos..gt_end]);
            sample_i += 1;
            pos = sample_end + 1;
        }
    } else {
        // General path: DS or GT at arbitrary index
        while sample_i < n_samples && pos < data_len {
            let sample_end = memchr::memchr(b'\t', &data[pos..])
                .map(|o| pos + o)
                .unwrap_or(data_len);

            let sample_slice = &data[pos..sample_end];

            out[sample_i] = if let Some(ds_i) = ds_idx {
                nth_subfield(sample_slice, ds_i)
                    .and_then(|s| std::str::from_utf8(s).ok())
                    .and_then(|s| s.parse::<f32>().ok())
                    .unwrap_or(f32::NAN)
            } else if let Some(gt_i) = gt_idx {
                nth_subfield(sample_slice, gt_i)
                    .map(gt_dosage_from_bytes)
                    .unwrap_or(f32::NAN)
            } else {
                f32::NAN
            };

            sample_i += 1;
            pos = sample_end + 1;
        }
    }

    // Pad with NAN if fewer samples than expected
    for i in sample_i..n_samples {
        out[i] = f32::NAN;
    }
}

/// Byte-level record parser for the streaming path.
///
/// Writes dosages directly into a caller-provided slice.
/// Returns variant metadata on success, or None for empty/short lines.
fn parse_record_into(
    line: &[u8],
    n_samples: usize,
    dosage_out: &mut [f32],
) -> Result<Option<(String, u32, String, Allele, Vec<Allele>, f64, String)>> {
    // Trim trailing whitespace
    let mut end = line.len();
    while end > 0 && (line[end - 1] == b'\n' || line[end - 1] == b'\r' || line[end - 1] == b' ')
    {
        end -= 1;
    }
    let line = &line[..end];
    if line.is_empty() {
        return Ok(None);
    }

    // ── Find first 9 tab positions using memchr ──────────────────
    let mut tabs = [0usize; 9];
    let mut pos = 0;
    for i in 0..9 {
        match memchr::memchr(b'\t', &line[pos..]) {
            Some(offset) => {
                tabs[i] = pos + offset;
                pos = pos + offset + 1;
            }
            None => return Ok(None), // fewer than 10 fields
        }
    }

    // ── Fixed fields (byte slices → minimal allocations) ─────────
    let chrom_bytes = &line[..tabs[0]];
    let pos_bytes = &line[tabs[0] + 1..tabs[1]];
    let id_bytes = &line[tabs[1] + 1..tabs[2]];
    let ref_bytes = &line[tabs[2] + 1..tabs[3]];
    let alt_bytes = &line[tabs[3] + 1..tabs[4]];
    let qual_bytes = &line[tabs[4] + 1..tabs[5]];
    let filter_bytes = &line[tabs[5] + 1..tabs[6]];
    // tabs[6]+1..tabs[7] is INFO – skipped
    let format_bytes = &line[tabs[7] + 1..tabs[8]];
    let sample_bytes = &line[tabs[8] + 1..];

    let chrom = string_from_bytes(chrom_bytes);
    let pos_val = parse_u32_fast(pos_bytes);

    let id = if id_bytes == b"." {
        format!(
            "{}:{}:{}:{}",
            chrom,
            string_from_bytes(pos_bytes),
            string_from_bytes(ref_bytes),
            string_from_bytes(alt_bytes),
        )
    } else {
        string_from_bytes(id_bytes)
    };

    let ref_str = std::str::from_utf8(ref_bytes).unwrap_or(".");
    let ref_allele = Allele::new(ref_str);

    let alt_str = std::str::from_utf8(alt_bytes).unwrap_or(".");
    let alt_alleles: Vec<Allele> = alt_str.split(',').map(Allele::new).collect();

    let qual = parse_f64_fast(qual_bytes);
    let filter = string_from_bytes(filter_bytes);

    // ── Parse FORMAT → find GT/DS indices ────────────────────────
    let gt_idx = find_format_index_bytes(format_bytes, b"GT");
    let ds_idx = find_format_index_bytes(format_bytes, b"DS");

    // ── Parse all samples directly into dosage_out ───────────────
    parse_samples_into(sample_bytes, n_samples, gt_idx, ds_idx, dosage_out);

    Ok(Some((
        chrom, pos_val, id, ref_allele, alt_alleles, qual, filter,
    )))
}

// ── Streaming VCF iterator (optimized) ───────────────────────────

struct VcfChunkedIter {
    reader: Box<dyn BufRead + Send>,
    n_samples: usize,
    chunk_size: usize,
    offset: usize,
    done: bool,
    /// Reusable line buffer – avoids per-line allocation
    line_buf: Vec<u8>,
    /// Reusable dosage buffer – avoids per-variant Vec<f32>
    dosage_buf: Vec<f32>,
}

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

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

        let mut chroms = Vec::with_capacity(self.chunk_size);
        let mut positions = Vec::with_capacity(self.chunk_size);
        let mut ids = Vec::with_capacity(self.chunk_size);
        let mut refs = Vec::with_capacity(self.chunk_size);
        let mut alts = Vec::with_capacity(self.chunk_size);
        let mut quals = Vec::with_capacity(self.chunk_size);
        let mut filters = Vec::with_capacity(self.chunk_size);

        // Build as (variants, samples) for cache-friendly row writes.
        // Transposed to API-expected (samples, variants) before returning.
        let mut matrix = Array2::<f32>::from_elem((self.chunk_size, self.n_samples), f32::NAN);

        let mut count = 0;

        while count < self.chunk_size {
            self.line_buf.clear();
            let n = match self.reader.read_until(b'\n', &mut self.line_buf) {
                Ok(n) => n,
                Err(e) => {
                    self.done = true;
                    return Some(Err(BunbunError::Io {
                        path: PathBuf::from("<vcf-stream>"),
                        source: e,
                    }));
                }
            };
            if n == 0 {
                self.done = true;
                break;
            }

            // Skip comment/header lines
            if !self.line_buf.is_empty() && self.line_buf[0] == b'#' {
                continue;
            }

            // parse_samples_into writes every position (and pads short
            // lines with NAN), so no need to pre-clear dosage_buf.

            match parse_record_into(&self.line_buf, self.n_samples, &mut self.dosage_buf) {
                Ok(Some((chrom, pos, id, ref_a, alt_a, qual, filt))) => {
                    // Write into row `count` — contiguous in memory
                    matrix.row_mut(count).assign(
                        &ndarray::ArrayView1::from(&self.dosage_buf[..self.n_samples]),
                    );
                    chroms.push(chrom);
                    positions.push(pos);
                    ids.push(id);
                    refs.push(ref_a);
                    alts.push(alt_a);
                    quals.push(qual);
                    filters.push(filt);
                    count += 1;
                }
                Ok(None) => continue,
                Err(e) => {
                    self.done = true;
                    return Some(Err(e));
                }
            }
        }

        if count == 0 {
            return None;
        }

        // Trim rows if fewer variants than chunk_size, then transpose
        // to the API-expected (samples, variants) layout.
        let trimmed = if count < self.chunk_size {
            matrix.slice(ndarray::s![..count, ..]).to_owned()
        } else {
            matrix
        };
        // .t() returns a transposed view; .to_owned() makes a contiguous
        // (samples, variants) array.  Single bulk copy — far cheaper than
        // scattered column writes during parsing.
        let final_matrix = trimmed.t().to_owned();

        let n_variants = count;
        let cm = vec![f64::NAN; n_variants];
        let variants = match VariantInfo::from_parts(
            chroms, positions, ids, refs, alts, cm, quals, filters,
        ) {
            Ok(v) => v,
            Err(e) => return Some(Err(e)),
        };

        let start = self.offset;
        self.offset += n_variants;

        // total_variants:
        //   0      → unknown (progress returns -1.0, is_last returns false)
        //   actual → stamped on the final chunk once EOF is reached
        let total = if self.done { self.offset } else { 0 };

        Some(Ok(GenotypeChunk {
            variants,
            dosages: DosageMatrix::new(final_matrix),
            offset: start,
            total_variants: total,
        }))
    }
}

// ── GenotypeReader impl ─────────────────────────────────────────

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

    fn read_samples(&self) -> Result<SampleInfo> {
        let mut reader = open_maybe_compressed(&self.path)?;
        let ids = Self::parse_header(&mut *reader)?;
        SampleInfo::from_ids(ids)
    }

    fn n_variants(&self) -> Result<usize> {
        Self::count_variants(&self.path)
    }

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

    fn read_chunked(&self, chunk_size: usize) -> Result<ChunkedReader> {
        // No pre-count pass — avoids decompressing the entire file just
        // to count lines.  total_variants is resolved lazily: 0 (unknown)
        // for intermediate chunks, correct total on the final chunk.

        let mut reader = open_maybe_compressed(&self.path)?;
        let sample_ids = Self::parse_header(&mut *reader)?;
        let n_samples = sample_ids.len();
        let samples = SampleInfo::from_ids(sample_ids)?;

        let iter = VcfChunkedIter {
            reader,
            n_samples,
            chunk_size,
            offset: 0,
            done: false,
            line_buf: Vec::with_capacity(n_samples * 8),
            dosage_buf: vec![f32::NAN; n_samples],
        };

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