use autark::mpera::programoutput::OutputValue;
use autark::onceframe::{OnceFrame, RealizedOnceFrame};
use autark::prelude::readers::transcend::TranscendReader;
use autark::prelude::readers::{csv::CsvReader, json::JsonReader, parquet::ParquetReader};
use pyo3::exceptions::PyKeyError;
use pyo3::prelude::*;
use pyo3::types::{PyDict, PyIterator, PyList, PyTuple};
use std::collections::HashMap;

use crate::arrow_py::record_batch_to_py_dict;
use crate::error::py_err;
use crate::program::PyProgram;
use crate::reader::{
    PyCsvReader, PyJsonReader, PyParquetReader, PyTranscendReader, ReaderSpec, reader_spec,
};
use crate::schema::PySchema;

fn onceframe_unavailable() -> PyErr {
    pyo3::exceptions::PyRuntimeError::new_err("OnceFrame unavailable")
}

fn frame_from_reader_spec(spec: ReaderSpec) -> PyResult<OnceFrame> {
    apply_reader_spec(None, spec)
}

fn with_reader_spec(frame: OnceFrame, spec: ReaderSpec) -> PyResult<OnceFrame> {
    apply_reader_spec(Some(frame), spec)
}

fn maybe_with_batch_size<T>(
    reader: T,
    batch_size: Option<usize>,
    apply: impl FnOnce(T, usize) -> T,
) -> T {
    match batch_size {
        Some(batch_size) => apply(reader, batch_size),
        None => reader,
    }
}

fn apply_reader_spec(frame: Option<OnceFrame>, spec: ReaderSpec) -> PyResult<OnceFrame> {
    Ok(match (frame, spec) {
        (None, ReaderSpec::Csv { uris, batch_size }) => OnceFrame::new(maybe_with_batch_size(
            CsvReader::from_uris(uris).map_err(py_err)?,
            batch_size,
            CsvReader::with_batch_size,
        )),
        (Some(frame), ReaderSpec::Csv { uris, batch_size }) => {
            frame.with_reader(maybe_with_batch_size(
                CsvReader::from_uris(uris).map_err(py_err)?,
                batch_size,
                CsvReader::with_batch_size,
            ))
        }
        (None, ReaderSpec::Parquet { uris, batch_size }) => OnceFrame::new(maybe_with_batch_size(
            ParquetReader::from_uris(uris).map_err(py_err)?,
            batch_size,
            ParquetReader::with_batch_size,
        )),
        (Some(frame), ReaderSpec::Parquet { uris, batch_size }) => {
            frame.with_reader(maybe_with_batch_size(
                ParquetReader::from_uris(uris).map_err(py_err)?,
                batch_size,
                ParquetReader::with_batch_size,
            ))
        }
        (None, ReaderSpec::JsonUri { uri, batch_size }) => OnceFrame::new(maybe_with_batch_size(
            JsonReader::new(&uri).map_err(py_err)?,
            batch_size,
            JsonReader::with_batch_size,
        )),
        (Some(frame), ReaderSpec::JsonUri { uri, batch_size }) => {
            frame.with_reader(maybe_with_batch_size(
                JsonReader::new(&uri).map_err(py_err)?,
                batch_size,
                JsonReader::with_batch_size,
            ))
        }
        (None, ReaderSpec::JsonInline { data, batch_size }) => {
            OnceFrame::new(maybe_with_batch_size(
                JsonReader::from_json_string(data).map_err(py_err)?,
                batch_size,
                JsonReader::with_batch_size,
            ))
        }
        (Some(frame), ReaderSpec::JsonInline { data, batch_size }) => {
            frame.with_reader(maybe_with_batch_size(
                JsonReader::from_json_string(data).map_err(py_err)?,
                batch_size,
                JsonReader::with_batch_size,
            ))
        }
        (
            None,
            ReaderSpec::Transcend {
                base_url,
                api_key,
                dataset,
                max_batches,
                latest,
            },
        ) => {
            let reader = transcend_reader(base_url, api_key, dataset, max_batches, latest)
                .map_err(py_err)?;
            OnceFrame::new(reader)
        }
        (
            Some(frame),
            ReaderSpec::Transcend {
                base_url,
                api_key,
                dataset,
                max_batches,
                latest,
            },
        ) => {
            let reader = transcend_reader(base_url, api_key, dataset, max_batches, latest)
                .map_err(py_err)?;
            frame.with_reader(reader)
        }
    })
}

fn transcend_reader(
    base_url: String,
    api_key: String,
    dataset: String,
    max_batches: Option<usize>,
    latest: Option<bool>,
) -> Result<TranscendReader, autark::prelude::Error> {
    let reader = TranscendReader::new(&base_url, &api_key, &dataset)?;
    let reader = match max_batches {
        Some(max_batches) => reader.with_max_batches(max_batches),
        None => reader,
    };
    Ok(match latest {
        Some(latest) => reader.with_latest(latest),
        None => reader,
    })
}

fn is_reader_object(item: &Bound<'_, PyAny>) -> bool {
    item.extract::<PyRef<'_, PyCsvReader>>().is_ok()
        || item.extract::<PyRef<'_, PyParquetReader>>().is_ok()
        || item.extract::<PyRef<'_, PyJsonReader>>().is_ok()
        || item.extract::<PyRef<'_, PyTranscendReader>>().is_ok()
}

fn frame_from_reader_input(reader: &Bound<'_, PyAny>) -> PyResult<OnceFrame> {
    if let Some(specs) = reader_specs_from_sequence(reader) {
        return frame_from_reader_specs(specs);
    }
    frame_from_reader_spec(reader_spec(reader)?)
}

fn sequence_items<'py>(reader: &Bound<'py, PyAny>) -> Option<Vec<Bound<'py, PyAny>>> {
    if let Ok(readers) = reader.downcast::<PyList>() {
        return Some(readers.iter().collect());
    }
    if let Ok(readers) = reader.downcast::<PyTuple>() {
        return Some(readers.iter().collect());
    }
    None
}

fn reader_specs_from_sequence(reader: &Bound<'_, PyAny>) -> Option<Vec<ReaderSpec>> {
    let items = sequence_items(reader)?;
    if items.is_empty() || !items.iter().all(is_reader_object) {
        return None;
    }
    Some(
        items
            .into_iter()
            .map(|item| reader_spec(&item))
            .collect::<PyResult<Vec<_>>>()
            .ok()?,
    )
}

fn frame_from_reader_specs(specs: Vec<ReaderSpec>) -> PyResult<OnceFrame> {
    let mut specs = specs.into_iter();
    let first = specs.next().ok_or_else(reader_spec_error)?;
    specs.try_fold(frame_from_reader_spec(first)?, with_reader_spec)
}

fn reader_spec_error() -> PyErr {
    pyo3::exceptions::PyTypeError::new_err("Expected a reader or sequence of readers")
}

fn frame_ref(frame: &Option<OnceFrame>) -> PyResult<&OnceFrame> {
    frame.as_ref().ok_or_else(onceframe_unavailable)
}

fn frame_mut(frame: &mut Option<OnceFrame>) -> PyResult<&mut OnceFrame> {
    frame.as_mut().ok_or_else(onceframe_unavailable)
}

fn take_frame(frame: &mut Option<OnceFrame>) -> PyResult<OnceFrame> {
    frame.take().ok_or_else(onceframe_unavailable)
}

fn output_value_to_py(py: Python<'_>, value: &OutputValue) -> PyResult<Py<PyAny>> {
    match value {
        OutputValue::Frame(batch) => Ok(record_batch_to_py_dict(py, batch)?.unbind().into_any()),
        OutputValue::Groups(groups) => {
            let out = PyDict::new(py);
            for (key, batch) in groups {
                out.set_item(key, record_batch_to_py_dict(py, batch)?)?;
            }
            Ok(out.unbind().into_any())
        }
    }
}

fn realized_batch(
    py: Python<'_>,
    realized: &RealizedOnceFrame,
    key: &str,
) -> PyResult<Option<Py<PyAny>>> {
    realized
        .output()
        .0
        .get(key)
        .map(|value| output_value_to_py(py, value))
        .transpose()
}

fn realized_items(
    py: Python<'_>,
    realized: &RealizedOnceFrame,
) -> PyResult<Vec<(String, Py<PyAny>)>> {
    realized
        .output()
        .0
        .iter()
        .map(|(key, value)| output_value_to_py(py, value).map(|obj| (key.clone(), obj)))
        .collect()
}

fn realized_values(py: Python<'_>, realized: &RealizedOnceFrame) -> PyResult<Vec<Py<PyAny>>> {
    realized_items(py, realized).map(|items| items.into_iter().map(|(_, value)| value).collect())
}

#[pyclass(name = "RealizedOnceFrame", unsendable)]
pub(crate) struct PyRealizedOnceFrame {
    inner: RealizedOnceFrame,
}

impl PyRealizedOnceFrame {
    fn new(inner: RealizedOnceFrame) -> PyRealizedOnceFrame {
        PyRealizedOnceFrame { inner }
    }
}

#[pymethods]
impl PyRealizedOnceFrame {
    fn __repr__(&self) -> PyResult<String> {
        Ok(self.inner.to_string())
    }

    fn __str__(&self) -> PyResult<String> {
        Ok(self.inner.to_string())
    }

    #[pyo3(signature = (key, default = None))]
    fn get(&self, py: Python<'_>, key: &str, default: Option<Py<PyAny>>) -> PyResult<Py<PyAny>> {
        if let Some(batch) = realized_batch(py, &self.inner, key)? {
            return Ok(batch.into_any());
        }
        Ok(default.unwrap_or_else(|| py.None()))
    }

    fn keys(&self) -> Vec<String> {
        self.inner.output().0.keys().cloned().collect()
    }

    fn values(&self, py: Python<'_>) -> PyResult<Vec<Py<PyAny>>> {
        realized_values(py, &self.inner)
    }

    fn items(&self, py: Python<'_>) -> PyResult<Vec<(String, Py<PyAny>)>> {
        realized_items(py, &self.inner)
    }

    fn __len__(&self) -> usize {
        self.inner.output().0.len()
    }

    fn __contains__(&self, key: &str) -> bool {
        self.inner.output().0.contains_key(key)
    }

    fn __getitem__(&self, py: Python<'_>, key: &str) -> PyResult<Py<PyAny>> {
        realized_batch(py, &self.inner, key)?.ok_or_else(|| PyKeyError::new_err(key.to_string()))
    }

    fn __iter__<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyIterator>> {
        let keys: Vec<String> = self.inner.output().0.keys().cloned().collect();
        PyIterator::from_object(&keys.into_pyobject(py)?)
    }
}

#[pyclass(name = "OnceFrame", unsendable)]
pub(crate) struct PyOnceFrame {
    frame: Option<OnceFrame>,
}

#[pymethods]
impl PyOnceFrame {
    #[new]
    fn new(reader: &Bound<'_, PyAny>) -> PyResult<PyOnceFrame> {
        let frame = frame_from_reader_input(reader)?;
        Ok(PyOnceFrame { frame: Some(frame) })
    }

    fn with_reader<'py>(
        mut slf: PyRefMut<'py, Self>,
        reader: &Bound<'_, PyAny>,
    ) -> PyResult<PyRefMut<'py, Self>> {
        let spec = reader_spec(reader)?;
        let frame = take_frame(&mut slf.frame)?;
        slf.frame = Some(with_reader_spec(frame, spec)?);
        Ok(slf)
    }

    fn __repr__(&self) -> PyResult<String> {
        frame_ref(&self.frame)?.describe(None).map_err(py_err)
    }

    fn __str__(&self) -> PyResult<String> {
        self.__repr__()
    }

    #[pyo3(signature = (index = None))]
    fn schema(&self, index: Option<usize>) -> PyResult<PySchema> {
        Ok(PySchema::new(
            frame_ref(&self.frame)?.schema(index).map_err(py_err)?,
        ))
    }

    #[pyo3(signature = (columns, index = None))]
    fn schema_of_columns(&self, columns: Vec<String>, index: Option<usize>) -> PyResult<PySchema> {
        let columns = columns.iter().map(String::as_str).collect::<Vec<_>>();
        Ok(PySchema::new(
            frame_ref(&self.frame)?
                .schema_of_columns(index, &columns)
                .map_err(py_err)?,
        ))
    }

    #[pyo3(signature = (index = None))]
    fn describe(&self, index: Option<usize>) -> PyResult<String> {
        frame_ref(&self.frame)?.describe(index).map_err(py_err)
    }

    #[pyo3(signature = (program, rename = None))]
    fn realize(
        &mut self,
        program: PyRef<'_, PyProgram>,
        rename: Option<Vec<HashMap<String, String>>>,
    ) -> PyResult<PyRealizedOnceFrame> {
        let realized = frame_mut(&mut self.frame)?
            .realize(program.inner_clone(), rename)
            .map_err(py_err)?;
        Ok(PyRealizedOnceFrame::new(realized))
    }
}
