use crate::inline_json::{inline_json_from_python_data, is_row_sequence};
use pyo3::exceptions::PyValueError;
use pyo3::prelude::*;
use pyo3::types::{PyDict, PyList, PyTuple};

#[pyclass(name = "CsvReader")]
#[derive(Clone)]
pub(crate) struct PyCsvReader {
    pub(crate) uris: Vec<String>,
    pub(crate) batch_size: Option<usize>,
}

#[pymethods]
impl PyCsvReader {
    #[new]
    #[pyo3(signature = (uri, batch_size = None))]
    fn new(uri: &Bound<'_, PyAny>, batch_size: Option<usize>) -> PyResult<PyCsvReader> {
        Ok(PyCsvReader {
            uris: path_list(uri)?,
            batch_size,
        })
    }
}

#[pyclass(name = "ParquetReader")]
#[derive(Clone)]
pub(crate) struct PyParquetReader {
    pub(crate) uris: Vec<String>,
    pub(crate) batch_size: Option<usize>,
}

#[pymethods]
impl PyParquetReader {
    #[new]
    #[pyo3(signature = (uri, batch_size = None))]
    fn new(uri: &Bound<'_, PyAny>, batch_size: Option<usize>) -> PyResult<PyParquetReader> {
        Ok(PyParquetReader {
            uris: path_list(uri)?,
            batch_size,
        })
    }
}

#[pyclass(name = "JsonReader")]
#[derive(Clone)]
pub(crate) struct PyJsonReader {
    pub(crate) source: PyJsonSource,
    pub(crate) batch_size: Option<usize>,
}

#[derive(Clone)]
pub(crate) enum PyJsonSource {
    Uri(String),
    Inline(String),
}

#[pymethods]
impl PyJsonReader {
    #[new]
    #[pyo3(signature = (source, batch_size = None))]
    fn new(
        py: Python<'_>,
        source: &Bound<'_, PyAny>,
        batch_size: Option<usize>,
    ) -> PyResult<PyJsonReader> {
        if let Ok(uri) = source.extract::<String>() {
            return Ok(PyJsonReader {
                source: PyJsonSource::Uri(uri),
                batch_size,
            });
        }

        Ok(PyJsonReader {
            source: PyJsonSource::Inline(inline_json_from_python_data(py, source)?),
            batch_size,
        })
    }
}

#[pyclass(name = "TranscendReader")]
#[derive(Clone)]
pub(crate) struct PyTranscendReader {
    pub(crate) base_url: String,
    pub(crate) api_key: String,
    pub(crate) dataset: String,
    pub(crate) max_batches: Option<usize>,
    pub(crate) latest: Option<bool>,
}

#[pymethods]
impl PyTranscendReader {
    #[new]
    #[pyo3(signature = (base_url, api_key, dataset, max_batches = None, latest = None))]
    fn new(
        base_url: String,
        api_key: String,
        dataset: String,
        max_batches: Option<usize>,
        latest: Option<bool>,
    ) -> PyResult<PyTranscendReader> {
        Ok(PyTranscendReader {
            base_url,
            api_key,
            dataset,
            max_batches,
            latest,
        })
    }
}

pub(crate) enum ReaderSpec {
    Csv {
        uris: Vec<String>,
        batch_size: Option<usize>,
    },
    Parquet {
        uris: Vec<String>,
        batch_size: Option<usize>,
    },
    JsonUri {
        uri: String,
        batch_size: Option<usize>,
    },
    JsonInline {
        data: String,
        batch_size: Option<usize>,
    },
    Transcend {
        base_url: String,
        api_key: String,
        dataset: String,
        max_batches: Option<usize>,
        latest: Option<bool>,
    },
}

pub(crate) fn reader_spec(reader: &Bound<'_, PyAny>) -> PyResult<ReaderSpec> {
    if let Ok(csv) = reader.extract::<PyRef<'_, PyCsvReader>>() {
        return Ok(ReaderSpec::Csv {
            uris: csv.uris.clone(),
            batch_size: csv.batch_size,
        });
    }
    if let Ok(parquet) = reader.extract::<PyRef<'_, PyParquetReader>>() {
        return Ok(ReaderSpec::Parquet {
            uris: parquet.uris.clone(),
            batch_size: parquet.batch_size,
        });
    }
    if let Ok(json) = reader.extract::<PyRef<'_, PyJsonReader>>() {
        return Ok(match &json.source {
            PyJsonSource::Uri(uri) => ReaderSpec::JsonUri {
                uri: uri.clone(),
                batch_size: json.batch_size,
            },
            PyJsonSource::Inline(data) => ReaderSpec::JsonInline {
                data: data.clone(),
                batch_size: json.batch_size,
            },
        });
    }
    if let Ok(transcend) = reader.extract::<PyRef<'_, PyTranscendReader>>() {
        return Ok(ReaderSpec::Transcend {
            base_url: transcend.base_url.clone(),
            api_key: transcend.api_key.clone(),
            dataset: transcend.dataset.clone(),
            max_batches: transcend.max_batches,
            latest: transcend.latest,
        });
    }

    if reader.is_instance_of::<PyDict>() || is_row_sequence(reader) {
        return Ok(ReaderSpec::JsonInline {
            data: inline_json_from_python_data(reader.py(), reader)?,
            batch_size: None,
        });
    }

    Err(PyValueError::new_err(
        "expected CsvReader, ParquetReader, JsonReader, TranscendReader, or JSON data (dict/list/tuple)",
    ))
}

fn path_list(value: &Bound<'_, PyAny>) -> PyResult<Vec<String>> {
    if let Ok(uri) = value.extract::<String>() {
        return Ok(vec![uri]);
    }
    if let Ok(paths) = value.downcast::<PyList>() {
        return paths
            .iter()
            .map(|item| item.extract::<String>())
            .collect::<PyResult<Vec<_>>>()
            .and_then(non_empty_paths);
    }
    if let Ok(paths) = value.downcast::<PyTuple>() {
        return paths
            .iter()
            .map(|item| item.extract::<String>())
            .collect::<PyResult<Vec<_>>>()
            .and_then(non_empty_paths);
    }

    Err(PyValueError::new_err(
        "expected a path string or sequence of path strings",
    ))
}

fn non_empty_paths(paths: Vec<String>) -> PyResult<Vec<String>> {
    if paths.is_empty() {
        return Err(PyValueError::new_err("expected at least one path string"));
    }
    Ok(paths)
}
