use numpy::{PyArray2, ToPyArray};
use pyo3::prelude::*;
use pyo3::types::PyDict;
use pyo3_polars::PyDataFrame;
use std::path::PathBuf;
use std::sync::Mutex;

use crate::core::GenotypeData;
use crate::formats::{BgenReader, PlinkReader, Plink2Reader, VcfReader};
use crate::io::{ChunkedReader, FormatDetector, GenotypeReader};

// ── GenotypeData (full read) ─────────────────────────────────────

#[pyclass(name = "GenotypeData")]
pub struct PyGenotypeData {
    inner: GenotypeData,
}

#[pymethods]
impl PyGenotypeData {
    #[getter]
    fn samples(&self) -> PyDataFrame {
        PyDataFrame(self.inner.samples.df.clone())
    }

    #[getter]
    fn variants(&self) -> PyDataFrame {
        PyDataFrame(self.inner.variants.df.clone())
    }

    #[getter]
    fn dosages<'py>(&self, py: Python<'py>) -> Bound<'py, PyArray2<f32>> {
        self.inner.dosages.as_array().to_pyarray(py)
    }

    #[getter]
    fn shape(&self) -> (usize, usize) {
        self.inner.shape()
    }

    #[getter]
    fn n_samples(&self) -> usize {
        self.inner.n_samples()
    }

    #[getter]
    fn n_variants(&self) -> usize {
        self.inner.n_variants()
    }

    #[getter]
    fn source_format(&self) -> &str {
        self.inner.source_format
    }

    fn allele_frequencies<'py>(&self, py: Python<'py>) -> Bound<'py, numpy::PyArray1<f32>> {
        numpy::PyArray1::from_vec(py, self.inner.dosages.allele_frequencies())
    }

    fn missing_rates<'py>(&self, py: Python<'py>) -> Bound<'py, numpy::PyArray1<f32>> {
        numpy::PyArray1::from_vec(py, self.inner.dosages.missing_rates())
    }

    fn mean_impute(&mut self) {
        self.inner.dosages.mean_impute();
    }

    fn subset_samples(&self, ids: Vec<String>) -> PyResult<Self> {
        let id_refs: Vec<&str> = ids.iter().map(|s| s.as_str()).collect();
        let data = self.inner.subset_samples(&id_refs).map_err(to_pyerr)?;
        Ok(Self { inner: data })
    }

    fn region(&self, chrom: &str, start: u32, end: u32) -> PyResult<Self> {
        let data = self.inner.region(chrom, start, end).map_err(to_pyerr)?;
        Ok(Self { inner: data })
    }

    fn __repr__(&self) -> String {
        format!(
            "GenotypeData(n_samples={}, n_variants={}, format='{}')",
            self.inner.n_samples(),
            self.inner.n_variants(),
            self.inner.source_format,
        )
    }

    fn __len__(&self) -> usize {
        self.inner.n_variants()
    }
}

// ── GenotypeChunk ────────────────────────────────────────────────

#[pyclass(name = "GenotypeChunk")]
pub struct PyGenotypeChunk {
    inner: crate::core::GenotypeChunk,
}

#[pymethods]
impl PyGenotypeChunk {
    #[getter]
    fn variants(&self) -> PyDataFrame {
        PyDataFrame(self.inner.variants.df.clone())
    }

    #[getter]
    fn dosages<'py>(&self, py: Python<'py>) -> Bound<'py, PyArray2<f32>> {
        self.inner.dosages.as_array().to_pyarray(py)
    }

    #[getter]
    fn n_samples(&self) -> usize {
        self.inner.n_samples()
    }

    #[getter]
    fn n_variants(&self) -> usize {
        self.inner.n_variants()
    }

    #[getter]
    fn offset(&self) -> usize {
        self.inner.offset
    }

    #[getter]
    fn total_variants(&self) -> usize {
        self.inner.total_variants
    }

    #[getter]
    fn progress(&self) -> f64 {
        self.inner.progress()
    }

    #[getter]
    fn is_last(&self) -> bool {
        self.inner.is_last()
    }

    fn __repr__(&self) -> String {
        format!(
            "GenotypeChunk(n_variants={}, offset={}/{}, progress={:.1}%)",
            self.inner.n_variants(),
            self.inner.offset,
            self.inner.total_variants,
            self.inner.progress() * 100.0,
        )
    }
}

// ── GenotypeStream ───────────────────────────────────────────────

#[pyclass(name = "GenotypeStream")]
pub struct PyGenotypeStream {
    reader: Mutex<ChunkedReader>,
    samples_df: PyDataFrame,
}

#[pymethods]
impl PyGenotypeStream {
    #[getter]
    fn samples(&self) -> PyDataFrame {
        self.samples_df.clone()
    }

    #[getter]
    fn source_format(&self) -> &str {
        self.reader.lock().unwrap().source_format
    }

    fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
        slf
    }

    fn __next__(&self) -> PyResult<Option<PyGenotypeChunk>> {
        let mut reader = self.reader.lock().unwrap();
        match reader.next() {
            Some(Ok(chunk)) => Ok(Some(PyGenotypeChunk { inner: chunk })),
            Some(Err(e)) => Err(to_pyerr(e)),
            None => Ok(None),
        }
    }

    fn __repr__(&self) -> String {
        let reader = self.reader.lock().unwrap();
        format!(
            "GenotypeStream(format='{}', n_samples={})",
            reader.source_format,
            reader.samples.n_samples(),
        )
    }
}

// ── Top-level functions ──────────────────────────────────────────

#[pyfunction]
#[pyo3(signature = (path))]
fn read(path: &str) -> PyResult<PyGenotypeData> {
    let path = PathBuf::from(path);
    let format = FormatDetector::detect(&path).map_err(to_pyerr)?;
    let reader: Box<dyn GenotypeReader> = match format {
        crate::io::Format::Plink => Box::new(PlinkReader::new(&path)),
        crate::io::Format::Plink2 => Box::new(Plink2Reader::new(&path)),
        crate::io::Format::Vcf => Box::new(VcfReader::new(&path)),
        crate::io::Format::Bgen => Box::new(BgenReader::new(&path)),
    };
    let data = reader.read_all().map_err(to_pyerr)?;
    Ok(PyGenotypeData { inner: data })
}

#[pyfunction]
#[pyo3(signature = (path, chunk_size=5000))]
fn stream(path: &str, chunk_size: usize) -> PyResult<PyGenotypeStream> {
    let path = PathBuf::from(path);
    let format = FormatDetector::detect(&path).map_err(to_pyerr)?;
    let reader: Box<dyn GenotypeReader> = match format {
        crate::io::Format::Plink => Box::new(PlinkReader::new(&path)),
        crate::io::Format::Plink2 => Box::new(Plink2Reader::new(&path)),
        crate::io::Format::Vcf => Box::new(VcfReader::new(&path)),
        crate::io::Format::Bgen => Box::new(BgenReader::new(&path)),
    };
    let chunked = reader.read_chunked(chunk_size).map_err(to_pyerr)?;
    let samples_df = PyDataFrame(chunked.samples.df.clone());
    Ok(PyGenotypeStream {
        reader: Mutex::new(chunked),
        samples_df,
    })
}

#[pyfunction]
#[pyo3(signature = (path, chunk_size=5000))]
fn stream_plink(path: &str, chunk_size: usize) -> PyResult<PyGenotypeStream> {
    let reader = PlinkReader::new(path);
    let chunked = reader.read_chunked(chunk_size).map_err(to_pyerr)?;
    let samples_df = PyDataFrame(chunked.samples.df.clone());
    Ok(PyGenotypeStream {
        reader: Mutex::new(chunked),
        samples_df,
    })
}

#[pyfunction]
#[pyo3(signature = (path, chunk_size=5000))]
fn stream_vcf(path: &str, chunk_size: usize) -> PyResult<PyGenotypeStream> {
    let reader = VcfReader::new(path);
    let chunked = reader.read_chunked(chunk_size).map_err(to_pyerr)?;
    let samples_df = PyDataFrame(chunked.samples.df.clone());
    Ok(PyGenotypeStream {
        reader: Mutex::new(chunked),
        samples_df,
    })
}

#[pyfunction]
#[pyo3(signature = (path, chunk_size=5000))]
fn stream_plink2(path: &str, chunk_size: usize) -> PyResult<PyGenotypeStream> {
    let reader = Plink2Reader::new(path);
    let chunked = reader.read_chunked(chunk_size).map_err(to_pyerr)?;
    let samples_df = PyDataFrame(chunked.samples.df.clone());
    Ok(PyGenotypeStream {
        reader: Mutex::new(chunked),
        samples_df,
    })
}

#[pyfunction]
#[pyo3(signature = (path, chunk_size=5000, sample_path=None))]
fn stream_bgen(
    path: &str,
    chunk_size: usize,
    sample_path: Option<&str>,
) -> PyResult<PyGenotypeStream> {
    let mut reader = BgenReader::new(path);
    if let Some(sp) = sample_path {
        reader = reader.with_sample_file(sp);
    }
    let chunked = reader.read_chunked(chunk_size).map_err(to_pyerr)?;
    let samples_df = PyDataFrame(chunked.samples.df.clone());
    Ok(PyGenotypeStream {
        reader: Mutex::new(chunked),
        samples_df,
    })
}

#[pyfunction]
#[pyo3(signature = (path))]
fn read_plink(path: &str) -> PyResult<PyGenotypeData> {
    let reader = PlinkReader::new(path);
    let data = reader.read_all().map_err(to_pyerr)?;
    Ok(PyGenotypeData { inner: data })
}

#[pyfunction]
#[pyo3(signature = (path))]
fn read_vcf(path: &str) -> PyResult<PyGenotypeData> {
    let reader = VcfReader::new(path);
    let data = reader.read_all().map_err(to_pyerr)?;
    Ok(PyGenotypeData { inner: data })
}

#[pyfunction]
#[pyo3(signature = (path))]
fn read_plink2(path: &str) -> PyResult<PyGenotypeData> {
    let reader = Plink2Reader::new(path);
    let data = reader.read_all().map_err(to_pyerr)?;
    Ok(PyGenotypeData { inner: data })
}

#[pyfunction]
#[pyo3(signature = (path, sample_path=None))]
fn read_bgen(path: &str, sample_path: Option<&str>) -> PyResult<PyGenotypeData> {
    let mut reader = BgenReader::new(path);
    if let Some(sp) = sample_path {
        reader = reader.with_sample_file(sp);
    }
    let data = reader.read_all().map_err(to_pyerr)?;
    Ok(PyGenotypeData { inner: data })
}

#[pyfunction]
fn supported_formats(py: Python<'_>) -> PyResult<Bound<'_, PyDict>> {
    let dict = PyDict::new(py);
    dict.set_item("plink", vec![".bed", ".bim", ".fam"])?;
    dict.set_item("plink2", vec![".pgen", ".pvar", ".psam"])?;
    dict.set_item("vcf", vec![".vcf", ".vcf.gz"])?;
    dict.set_item("bgen", vec![".bgen", ".sample"])?;
    Ok(dict)
}

// ── Module registration ──────────────────────────────────────────

pub fn register_module(parent: &Bound<'_, PyModule>) -> PyResult<()> {
    parent.add_class::<PyGenotypeData>()?;
    parent.add_class::<PyGenotypeChunk>()?;
    parent.add_class::<PyGenotypeStream>()?;

    parent.add_function(wrap_pyfunction!(read, parent)?)?;
    parent.add_function(wrap_pyfunction!(read_plink, parent)?)?;
    parent.add_function(wrap_pyfunction!(read_vcf, parent)?)?;
    parent.add_function(wrap_pyfunction!(read_plink2, parent)?)?;
    parent.add_function(wrap_pyfunction!(read_bgen, parent)?)?;

    parent.add_function(wrap_pyfunction!(stream, parent)?)?;
    parent.add_function(wrap_pyfunction!(stream_plink, parent)?)?;
    parent.add_function(wrap_pyfunction!(stream_vcf, parent)?)?;
    parent.add_function(wrap_pyfunction!(stream_plink2, parent)?)?;
    parent.add_function(wrap_pyfunction!(stream_bgen, parent)?)?;

    parent.add_function(wrap_pyfunction!(supported_formats, parent)?)?;
    Ok(())
}

fn to_pyerr(e: crate::error::BunbunError) -> PyErr {
    pyo3::exceptions::PyIOError::new_err(e.to_string())
}
