use pyo3::exceptions::PyValueError;
use pyo3::prelude::*;
use pyo3::types::{PyDict, PyList, PyTuple};

pub(crate) fn is_row_sequence(value: &Bound<'_, PyAny>) -> bool {
    value.is_instance_of::<PyList>() || value.is_instance_of::<PyTuple>()
}

fn ndjson_from_rows(py: Python<'_>, source: &Bound<'_, PyAny>) -> PyResult<String> {
    let json = py.import("json")?;
    let items = source.try_iter()?.collect::<PyResult<Vec<_>>>()?;
    let mut out = String::new();

    for (i, item) in items.iter().enumerate() {
        let line: String = json.call_method1("dumps", (item,))?.extract()?;
        if i > 0 {
            out.push('\n');
        }
        out.push_str(&line);
    }

    Ok(out)
}

fn ndjson_from_column_oriented_dict(
    py: Python<'_>,
    dict: &Bound<'_, PyDict>,
) -> PyResult<Option<String>> {
    if dict.is_empty() {
        return Ok(None);
    }

    let mut columns: Vec<(String, Vec<Py<PyAny>>)> = Vec::new();
    let mut row_count: Option<usize> = None;

    for (key, value) in dict.iter() {
        if !is_row_sequence(&value) {
            return Ok(None);
        }

        let key = key.extract::<String>().map_err(|_| {
            PyValueError::new_err("column-oriented JSON object must use string keys")
        })?;

        let items = value
            .try_iter()?
            .map(|item| item.map(|obj| obj.unbind()))
            .collect::<PyResult<Vec<_>>>()?;

        if let Some(expected) = row_count {
            if items.len() != expected {
                return Err(PyValueError::new_err(
                    "column-oriented JSON object must use equal-length value sequences",
                ));
            }
        } else {
            row_count = Some(items.len());
        }

        columns.push((key, items));
    }

    let json = py.import("json")?;
    let mut out = String::new();

    for row_idx in 0..row_count.unwrap_or(0) {
        let row = PyDict::new(py);
        for (key, col) in &columns {
            row.set_item(key, col[row_idx].bind(py))?;
        }

        let line: String = json.call_method1("dumps", (row,))?.extract()?;
        if row_idx > 0 {
            out.push('\n');
        }
        out.push_str(&line);
    }

    Ok(Some(out))
}

pub(crate) fn inline_json_from_python_data(
    py: Python<'_>,
    source: &Bound<'_, PyAny>,
) -> PyResult<String> {
    if is_row_sequence(source) {
        return ndjson_from_rows(py, source);
    }

    if source.is_instance_of::<PyDict>() {
        let dict = source.cast::<PyDict>()?;

        if let Some(rows) = ndjson_from_column_oriented_dict(py, &dict)? {
            return Ok(rows);
        }

        let json = py.import("json")?;
        let line: String = json.call_method1("dumps", (source,))?.extract()?;
        return Ok(line);
    }

    Err(PyValueError::new_err(
        "JSON data must be a dict, list, or tuple",
    ))
}
