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)
    }

    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::with_capacity(n_samples);

        for (i, sample_field) in sample_block.split('\t').enumerate() {
            if i >= n_samples {
                break;
            }
            let subfields: Vec<&str> = sample_field.split(':').collect();

            let dosage = if let Some(ds_i) = ds_idx {
                subfields
                    .get(ds_i)
                    .and_then(|s| s.parse::<f32>().ok())
                    .unwrap_or(f32::NAN)
            } else if let Some(gt_i) = gt_idx {
                subfields
                    .get(gt_i)
                    .map(|gt| GenotypeValue::from_vcf_gt(gt).to_dosage())
                    .unwrap_or(f32::NAN)
            } else {
                f32::NAN
            };

            dosages.push(dosage);
        }

        dosages.resize(n_samples, f32::NAN);

        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)))
    }
}

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

/// Streaming VCF iterator
struct VcfChunkedIter {
    reader: Box<dyn BufRead + Send>,
    n_samples: usize,
    chunk_size: usize,
    offset: usize,
    total_variants: usize,
    done: bool,
}

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);
        let mut dosage_rows: Vec<Vec<f32>> = Vec::with_capacity(self.chunk_size);

        let mut line = String::new();
        let mut count = 0;

        while count < self.chunk_size {
            line.clear();
            let n = match self.reader.read_line(&mut line) {
                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;
            }

            match VcfReader::parse_record(&line, self.n_samples) {
                Ok(Some((chrom, pos, id, ref_a, alt_a, qual, filt, dosages))) => {
                    chroms.push(chrom);
                    positions.push(pos);
                    ids.push(id);
                    refs.push(ref_a);
                    alts.push(alt_a);
                    quals.push(qual);
                    filters.push(filt);
                    dosage_rows.push(dosages);
                    count += 1;
                }
                Ok(None) => continue,
                Err(e) => {
                    self.done = true;
                    return Some(Err(e));
                }
            }
        }

        if count == 0 {
            return None;
        }

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

        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;

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

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> {
        let total_variants = Self::count_variants(&self.path)?;

        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,
            total_variants,
            done: false,
        };

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