use pyo3::exceptions::{PyIOError, PyStopIteration, PyValueError};
use pyo3::prelude::*;
use pyo3_polars::PyDataFrame;
use rssa_core::{Column, RssaError, RssaFile};

/// Converts an `RssaError` into an appropriate Python exception.
fn to_py_err(err: RssaError) -> PyErr {
    match err {
        RssaError::Io(e) => PyIOError::new_err(e.to_string()),
        other => PyValueError::new_err(other.to_string()),
    }
}

/// Parses a list of column names into `rssa_core::Column`s, in the given order.
/// `None` means "all columns, in their canonical order".
fn resolve_columns(columns: Option<Vec<String>>) -> PyResult<Vec<Column>> {
    match columns {
        None => Ok(Column::ALL.to_vec()),
        Some(names) => names
            .iter()
            .map(|name| {
                Column::from_name(name).ok_or_else(|| {
                    PyValueError::new_err(format!(
                        "unknown track column {name:?}; valid columns are: {:?}",
                        Column::ALL.map(Column::name)
                    ))
                })
            })
            .collect(),
    }
}

/// All valid track column names, in their canonical order.
#[pyfunction]
fn column_names() -> Vec<&'static str> {
    Column::ALL.map(Column::name).to_vec()
}

/// Parses only the header of an RSSA file and returns its parameters as JSON.
///
/// A thin wrapper: the Python-side `FileParameters` dataclass is built from this JSON
/// so that the parsing logic (and its correctness guarantees) live in one place, in
/// Rust.
#[pyfunction]
fn parse_header_json(py: Python<'_>, path: String) -> PyResult<String> {
    py.detach(|| {
        let file = RssaFile::open(&path).map_err(to_py_err)?;
        file.parameters().to_json().map_err(to_py_err)
    })
}

/// Eagerly decodes track data from an RSSA file into a Polars DataFrame.
///
/// If `n_rows` is given, only the first `n_rows` tracks are decoded (the rest of the
/// file is never read), which keeps `.head(n)`-style queries on very large files
/// cheap.
#[pyfunction]
#[pyo3(signature = (path, columns=None, n_rows=None))]
fn read_tracks(
    py: Python<'_>,
    path: String,
    columns: Option<Vec<String>>,
    n_rows: Option<usize>,
) -> PyResult<PyDataFrame> {
    let columns = resolve_columns(columns)?;
    py.detach(|| {
        let file = RssaFile::open(&path).map_err(to_py_err)?;
        let df = file.read_tracks(&columns, n_rows).map_err(to_py_err)?;
        Ok(PyDataFrame(df))
    })
}

/// Streaming, batch-at-a-time reader over the track data of an RSSA file, exposed to
/// Python as an iterator of `polars.DataFrame` batches.
#[pyclass(module = "aci_rssa._aci_rssa")]
struct RssaBatchReader {
    file: RssaFile,
    batch_size: usize,
    columns: Vec<Column>,
    inner: Option<rssa_core::BatchReader>,
}

#[pymethods]
impl RssaBatchReader {
    #[new]
    #[pyo3(signature = (path, batch_size, columns=None))]
    fn new(
        py: Python<'_>,
        path: String,
        batch_size: usize,
        columns: Option<Vec<String>>,
    ) -> PyResult<Self> {
        let columns = resolve_columns(columns)?;
        py.detach(|| {
            let file = RssaFile::open(&path).map_err(to_py_err)?;
            Ok(Self {
                file,
                batch_size,
                columns,
                inner: None,
            })
        })
    }

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

    fn __next__(mut slf: PyRefMut<'_, Self>, py: Python<'_>) -> PyResult<PyDataFrame> {
        if slf.inner.is_none() {
            let batch_size = slf.batch_size;
            let columns = slf.columns.clone();
            slf.inner = Some(slf.file.batches(batch_size, columns));
        }
        let reader = slf.inner.as_mut().unwrap();
        let batch = py.detach(|| reader.next_batch()).map_err(to_py_err)?;
        match batch {
            Some(df) => Ok(PyDataFrame(df)),
            None => Err(PyStopIteration::new_err(())),
        }
    }
}

#[pymodule]
fn _aci_rssa(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(column_names, m)?)?;
    m.add_function(wrap_pyfunction!(parse_header_json, m)?)?;
    m.add_function(wrap_pyfunction!(read_tracks, m)?)?;
    m.add_class::<RssaBatchReader>()?;
    Ok(())
}
