use std::sync::Arc;

use pyo3::exceptions::PyIndexError;
use pyo3::exceptions::PyKeyError;
use pyo3::exceptions::PyNotImplementedError;
use pyo3::exceptions::PyValueError;
use pyo3::prelude::*;
use pyo3::types::PyAny;
use pyo3::types::PyBytes;
use pyo3::types::PyCapsule;
use pyo3::types::PyDict;
use pyo3::types::PyIterator;
use pyo3::types::PyList;
use pyo3::types::PyTuple;
use pyo3_arrow::ffi::ArrayIterator;
use pyo3_arrow::ffi::to_array_pycapsules;
use pyo3_arrow::ffi::to_stream_pycapsule;
use qroissant_arrow::HeterogeneousListMode;
use qroissant_arrow::IngestionError;
use qroissant_arrow::ProjectionOptions;
use qroissant_arrow::ingest_array;
use qroissant_arrow::ingest_record_batch;
use qroissant_arrow::ingest_record_batch_reader;
use qroissant_arrow::project;
use qroissant_arrow::project_table;
use qroissant_core::Atom as CoreAtom;
use qroissant_core::Dictionary as CoreDictionary;
use qroissant_core::List as CoreList;
use qroissant_core::Table as CoreTable;
use qroissant_core::Value as CoreValue;
use qroissant_core::Vector as CoreVector;
use qroissant_core::VectorData;

use crate::errors::to_py_err;
use crate::types::Attribute;
use crate::types::Compression;
use crate::types::Encoding;
use crate::types::MessageType;
use crate::types::Primitive;
use crate::types::Shape;
use crate::types::Type;

#[pyclass(subclass, module = "qroissant")]
#[derive(Clone, Debug)]
pub struct Value {
    inner: CoreValue,
    projection_opts: Arc<ProjectionOptions>,
}

impl Value {
    pub fn new(inner: CoreValue) -> Self {
        Self {
            inner,
            projection_opts: Arc::new(ProjectionOptions::default()),
        }
    }

    pub fn new_with_opts(inner: CoreValue, opts: Arc<ProjectionOptions>) -> Self {
        Self {
            inner,
            projection_opts: opts,
        }
    }

    pub fn inner(&self) -> &CoreValue {
        &self.inner
    }

    pub fn into_inner(self) -> CoreValue {
        self.inner
    }

    pub fn projection_opts(&self) -> &Arc<ProjectionOptions> {
        &self.projection_opts
    }
}

#[pymethods]
impl Value {
    #[getter]
    fn qtype(&self) -> Type {
        Type::from(self.inner.qtype())
    }

    #[getter]
    fn primitive(&self) -> Option<Primitive> {
        self.inner.qtype().primitive.map(Primitive::from)
    }

    #[getter]
    fn shape(&self) -> Shape {
        Shape::from(self.inner.qtype().shape)
    }

    #[getter]
    fn attribute(&self) -> Option<Attribute> {
        self.inner.qtype().attribute.map(Attribute::from)
    }

    #[pyo3(signature = (*, options=None, encoding=Encoding::LittleEndian, message_type=MessageType::Asynchronous, compression=Compression::Uncompressed))]
    fn serialize(
        &self,
        options: Option<&crate::types::EncodeOptions>,
        encoding: Encoding,
        message_type: MessageType,
        compression: Compression,
    ) -> PyResult<Py<PyBytes>> {
        let inner = self.inner.clone();
        let options_clone = options.cloned();
        Python::attach(|py| {
            let payload = py
                .detach(|| {
                    crate::serde::encode_core_value_bytes(
                        &inner,
                        options_clone.as_ref(),
                        encoding,
                        message_type,
                        compression,
                    )
                })
                .map_err(to_py_err)?;
            Ok(PyBytes::new(py, &payload).unbind())
        })
    }
}

#[pyclass(extends = Value, module = "qroissant")]
#[derive(Clone, Debug)]
pub struct Atom;

#[pymethods]
impl Atom {
    #[new]
    fn new(qtype: PyRef<'_, Type>, value: &Bound<'_, PyAny>) -> PyResult<(Self, Value)> {
        let core = atom_from_python(&qtype, value)?;
        Ok((Self, Value::new(CoreValue::Atom(core))))
    }

    fn as_py(slf: PyRef<'_, Self>, py: Python<'_>) -> PyResult<Py<PyAny>> {
        match slf.as_super().inner() {
            CoreValue::Atom(atom) => atom_to_python(py, atom),
            _ => unreachable!("Atom instances always hold q atoms"),
        }
    }

    #[getter]
    fn value(slf: PyRef<'_, Self>, py: Python<'_>) -> PyResult<Py<PyAny>> {
        Self::as_py(slf, py)
    }

    fn is_null(slf: PyRef<'_, Self>) -> bool {
        use qroissant_kernels::nulls::*;
        match slf.as_super().inner() {
            CoreValue::Atom(atom) => match atom {
                CoreAtom::Boolean(_)
                | CoreAtom::Guid(_)
                | CoreAtom::Byte(_)
                | CoreAtom::Char(_)
                | CoreAtom::Symbol(_) => false,
                CoreAtom::Short(v) => *v == Q_NULL_SHORT,
                CoreAtom::Int(v) => *v == Q_NULL_INT,
                CoreAtom::Long(v) => *v == Q_NULL_LONG,
                CoreAtom::Real(v) => v.is_nan(),
                CoreAtom::Float(v) => v.is_nan(),
                CoreAtom::Timestamp(v) => *v == Q_NULL_TIMESTAMP,
                CoreAtom::Month(v) => *v == Q_NULL_MONTH,
                CoreAtom::Date(v) => *v == Q_NULL_DATE,
                CoreAtom::Datetime(v) => v.is_nan(),
                CoreAtom::Timespan(v) => *v == Q_NULL_TIMESPAN,
                CoreAtom::Minute(v) => *v == Q_NULL_MINUTE,
                CoreAtom::Second(v) => *v == Q_NULL_SECOND,
                CoreAtom::Time(v) => *v == Q_NULL_TIME,
            },
            _ => unreachable!("Atom instances always hold q atoms"),
        }
    }

    fn is_infinite(slf: PyRef<'_, Self>) -> bool {
        use qroissant_kernels::nulls::*;
        match slf.as_super().inner() {
            CoreValue::Atom(atom) => match atom {
                CoreAtom::Boolean(_)
                | CoreAtom::Guid(_)
                | CoreAtom::Byte(_)
                | CoreAtom::Char(_)
                | CoreAtom::Symbol(_) => false,
                CoreAtom::Short(v) => *v == Q_INF_SHORT || *v == Q_NINF_SHORT,
                CoreAtom::Int(v) => *v == Q_INF_INT || *v == Q_NINF_INT,
                CoreAtom::Long(v) => *v == Q_INF_LONG || *v == Q_NINF_LONG,
                CoreAtom::Real(v) => v.is_infinite(),
                CoreAtom::Float(v) => v.is_infinite(),
                CoreAtom::Timestamp(v) => *v == Q_INF_TIMESTAMP || *v == Q_NINF_TIMESTAMP,
                CoreAtom::Month(v) => *v == Q_INF_MONTH || *v == Q_NINF_MONTH,
                CoreAtom::Date(v) => *v == Q_INF_DATE || *v == Q_NINF_DATE,
                CoreAtom::Datetime(v) => v.is_infinite(),
                CoreAtom::Timespan(v) => *v == Q_INF_TIMESPAN || *v == Q_NINF_TIMESPAN,
                CoreAtom::Minute(v) => *v == Q_INF_MINUTE || *v == Q_NINF_MINUTE,
                CoreAtom::Second(v) => *v == Q_INF_SECOND || *v == Q_NINF_SECOND,
                CoreAtom::Time(v) => *v == Q_INF_TIME || *v == Q_NINF_TIME,
            },
            _ => unreachable!("Atom instances always hold q atoms"),
        }
    }

    #[pyo3(signature = (requested_schema=None))]
    fn __arrow_c_array__(
        slf: PyRef<'_, Self>,
        py: Python<'_>,
        requested_schema: Option<Bound<'_, PyAny>>,
    ) -> PyResult<Py<PyTuple>> {
        let schema_capsule: Option<Bound<'_, PyCapsule>> = requested_schema
            .map(|s| s.downcast_into::<PyCapsule>())
            .transpose()?;
        let opts = effective_export_opts(py, slf.as_super().projection_opts());
        let export = project(slf.as_super().inner(), &opts)
            .map_err(|e| PyNotImplementedError::new_err(e.to_string()))?;
        let capsules =
            to_array_pycapsules(py, export.field, export.array.as_ref(), schema_capsule)?;
        Ok(capsules.unbind())
    }

    fn __repr__(slf: PyRef<'_, Self>) -> String {
        match slf.as_super().inner() {
            CoreValue::Atom(atom) => crate::repr::format_atom(atom),
            _ => unreachable!("Atom instances always hold q atoms"),
        }
    }

    fn __str__(slf: PyRef<'_, Self>) -> String {
        Self::__repr__(slf)
    }
}

#[pyclass(extends = Value, module = "qroissant")]
#[derive(Clone, Debug)]
pub struct Vector;

#[pymethods]
impl Vector {
    #[new]
    fn new(qtype: PyRef<'_, Type>, values: Option<&Bound<'_, PyAny>>) -> PyResult<(Self, Value)> {
        let core = if let Some(values) = values {
            vector_from_python(&qtype, values)?
        } else {
            let empty = PyList::empty(qtype.py());
            vector_from_python(&qtype, empty.as_any())?
        };
        Ok((Self, Value::new(CoreValue::Vector(core))))
    }

    fn __len__(slf: PyRef<'_, Self>) -> usize {
        match slf.as_super().inner() {
            CoreValue::Vector(vector) => vector.len(),
            _ => unreachable!("Vector instances always hold q vectors"),
        }
    }

    fn __iter__(slf: PyRef<'_, Self>, py: Python<'_>) -> PyResult<Py<PyAny>> {
        let list = Self::to_list(slf, py)?;
        let iter = PyIterator::from_object(list.bind(py).as_any())?;
        Ok(iter.into_any().unbind())
    }

    fn __getitem__(slf: PyRef<'_, Self>, py: Python<'_>, index: isize) -> PyResult<Py<PyAny>> {
        let vector = match slf.as_super().inner() {
            CoreValue::Vector(vector) => vector,
            _ => unreachable!("Vector instances always hold q vectors"),
        };
        let index = normalize_index(index, vector.len())?;
        vector_item_to_python(py, vector, index)
    }

    fn to_list(slf: PyRef<'_, Self>, py: Python<'_>) -> PyResult<Py<PyList>> {
        let vector = match slf.as_super().inner() {
            CoreValue::Vector(vector) => vector,
            _ => unreachable!("Vector instances always hold q vectors"),
        };
        vector_to_pylist(py, vector)
    }

    #[pyo3(signature = (requested_schema=None))]
    fn __arrow_c_array__(
        slf: PyRef<'_, Self>,
        py: Python<'_>,
        requested_schema: Option<Bound<'_, PyAny>>,
    ) -> PyResult<Py<PyTuple>> {
        let schema_capsule: Option<Bound<'_, PyCapsule>> = requested_schema
            .map(|s| s.downcast_into::<PyCapsule>())
            .transpose()?;
        let opts = effective_export_opts(py, slf.as_super().projection_opts());
        let export = project(slf.as_super().inner(), &opts)
            .map_err(|e| PyNotImplementedError::new_err(e.to_string()))?;
        let capsules =
            to_array_pycapsules(py, export.field, export.array.as_ref(), schema_capsule)?;
        Ok(capsules.unbind())
    }

    fn __repr__(slf: PyRef<'_, Self>) -> String {
        match slf.as_super().inner() {
            CoreValue::Vector(vector) => crate::repr::format_vector(vector),
            _ => unreachable!("Vector instances always hold q vectors"),
        }
    }

    fn __str__(slf: PyRef<'_, Self>) -> String {
        Self::__repr__(slf)
    }
}

#[pyclass(extends = Value, module = "qroissant")]
#[derive(Clone, Debug)]
pub struct List;

#[pymethods]
impl List {
    #[new]
    fn new(qtype: PyRef<'_, Type>, values: Option<&Bound<'_, PyAny>>) -> PyResult<(Self, Value)> {
        let core = if let Some(values) = values {
            list_from_python(&qtype, values)?
        } else {
            let empty = PyList::empty(qtype.py());
            list_from_python(&qtype, empty.as_any())?
        };
        Ok((Self, Value::new(CoreValue::List(core))))
    }

    fn __len__(slf: PyRef<'_, Self>) -> usize {
        match slf.as_super().inner() {
            CoreValue::List(list) => list.len(),
            _ => unreachable!("List instances always hold q lists"),
        }
    }

    fn __iter__(slf: PyRef<'_, Self>, py: Python<'_>) -> PyResult<Py<PyAny>> {
        let list = Self::to_list(slf, py)?;
        let iter = PyIterator::from_object(list.bind(py).as_any())?;
        Ok(iter.into_any().unbind())
    }

    fn __getitem__(slf: PyRef<'_, Self>, py: Python<'_>, index: isize) -> PyResult<Py<PyAny>> {
        let list = match slf.as_super().inner() {
            CoreValue::List(list) => list,
            _ => unreachable!("List instances always hold q lists"),
        };
        let index = normalize_index(index, list.len())?;
        core_value_to_python(py, list.values()[index].clone())
    }

    fn to_list(slf: PyRef<'_, Self>, py: Python<'_>) -> PyResult<Py<PyList>> {
        let list = match slf.as_super().inner() {
            CoreValue::List(list) => list,
            _ => unreachable!("List instances always hold q lists"),
        };
        let mut values = Vec::with_capacity(list.len());
        for value in list.values() {
            values.push(core_value_to_python(py, value.clone())?);
        }
        Ok(PyList::new(py, values)?.unbind())
    }

    #[pyo3(signature = (requested_schema=None))]
    fn __arrow_c_array__(
        slf: PyRef<'_, Self>,
        py: Python<'_>,
        requested_schema: Option<Bound<'_, PyAny>>,
    ) -> PyResult<Py<PyTuple>> {
        let schema_capsule: Option<Bound<'_, PyCapsule>> = requested_schema
            .map(|s| s.downcast_into::<PyCapsule>())
            .transpose()?;
        let opts = effective_export_opts(py, slf.as_super().projection_opts());
        let export = project(slf.as_super().inner(), &opts)
            .map_err(|e| PyNotImplementedError::new_err(e.to_string()))?;
        let capsules =
            to_array_pycapsules(py, export.field, export.array.as_ref(), schema_capsule)?;
        Ok(capsules.unbind())
    }

    fn __repr__(slf: PyRef<'_, Self>) -> String {
        match slf.as_super().inner() {
            CoreValue::List(list) => crate::repr::format_list(list),
            _ => unreachable!("List instances always hold q lists"),
        }
    }

    fn __str__(slf: PyRef<'_, Self>) -> String {
        Self::__repr__(slf)
    }
}

#[pyclass(extends = Value, module = "qroissant")]
#[derive(Clone, Debug)]
pub struct Dictionary;

#[pymethods]
impl Dictionary {
    #[new]
    fn new(
        qtype: PyRef<'_, Type>,
        keys: &Bound<'_, PyAny>,
        values: &Bound<'_, PyAny>,
    ) -> PyResult<(Self, Value)> {
        let core = dictionary_from_python(&qtype, keys, values)?;
        Ok((Self, Value::new(CoreValue::Dictionary(core))))
    }

    #[getter]
    fn keys(slf: PyRef<'_, Self>, py: Python<'_>) -> PyResult<Py<PyAny>> {
        match slf.as_super().inner() {
            CoreValue::Dictionary(dictionary) => {
                core_value_to_python(py, dictionary.keys().clone())
            }
            _ => unreachable!("Dictionary instances always hold q dictionaries"),
        }
    }

    #[getter]
    fn values(slf: PyRef<'_, Self>, py: Python<'_>) -> PyResult<Py<PyAny>> {
        match slf.as_super().inner() {
            CoreValue::Dictionary(dictionary) => {
                core_value_to_python(py, dictionary.values().clone())
            }
            _ => unreachable!("Dictionary instances always hold q dictionaries"),
        }
    }

    fn __len__(slf: PyRef<'_, Self>) -> usize {
        match slf.as_super().inner() {
            CoreValue::Dictionary(dictionary) => dictionary.len(),
            _ => unreachable!("Dictionary instances always hold q dictionaries"),
        }
    }

    #[pyo3(signature = (requested_schema=None))]
    fn __arrow_c_array__(
        slf: PyRef<'_, Self>,
        py: Python<'_>,
        requested_schema: Option<Bound<'_, PyAny>>,
    ) -> PyResult<Py<PyTuple>> {
        let schema_capsule: Option<Bound<'_, PyCapsule>> = requested_schema
            .map(|s| s.downcast_into::<PyCapsule>())
            .transpose()?;
        let opts = effective_export_opts(py, slf.as_super().projection_opts());
        let export = project(slf.as_super().inner(), &opts)
            .map_err(|e| PyNotImplementedError::new_err(e.to_string()))?;
        let capsules =
            to_array_pycapsules(py, export.field, export.array.as_ref(), schema_capsule)?;
        Ok(capsules.unbind())
    }

    fn __repr__(slf: PyRef<'_, Self>) -> String {
        match slf.as_super().inner() {
            CoreValue::Dictionary(dict) => crate::repr::format_dictionary(dict),
            _ => unreachable!("Dictionary instances always hold q dictionaries"),
        }
    }

    fn __str__(slf: PyRef<'_, Self>) -> String {
        Self::__repr__(slf)
    }
}

#[pyclass(extends = Value, module = "qroissant")]
#[derive(Clone, Debug)]
pub struct Table;

#[pymethods]
impl Table {
    #[new]
    fn new(qtype: PyRef<'_, Type>, columns: Option<&Bound<'_, PyAny>>) -> PyResult<(Self, Value)> {
        let core = if let Some(columns) = columns {
            table_from_python(&qtype, columns)?
        } else {
            let empty = PyDict::new(qtype.py());
            table_from_python(&qtype, empty.as_any())?
        };
        Ok((Self, Value::new(CoreValue::Table(core))))
    }

    #[getter]
    fn columns(slf: PyRef<'_, Self>) -> PyResult<Vec<String>> {
        match slf.as_super().inner() {
            CoreValue::Table(table) => table
                .column_names()
                .iter()
                .map(|name| {
                    String::from_utf8(name.to_vec()).map_err(|_| {
                        PyValueError::new_err("q table column names must be valid UTF-8 for now")
                    })
                })
                .collect(),
            _ => unreachable!("Table instances always hold q tables"),
        }
    }

    #[getter]
    fn num_rows(slf: PyRef<'_, Self>) -> usize {
        match slf.as_super().inner() {
            CoreValue::Table(table) => table.len(),
            _ => unreachable!("Table instances always hold q tables"),
        }
    }

    #[getter]
    fn num_columns(slf: PyRef<'_, Self>) -> usize {
        match slf.as_super().inner() {
            CoreValue::Table(table) => table.num_columns(),
            _ => unreachable!("Table instances always hold q tables"),
        }
    }

    fn column(slf: PyRef<'_, Self>, py: Python<'_>, name: &str) -> PyResult<Py<PyAny>> {
        match slf.as_super().inner() {
            CoreValue::Table(table) => {
                let needle = name.as_bytes();
                for (idx, candidate) in table.column_names().iter().enumerate() {
                    if candidate.as_ref() == needle {
                        return core_value_to_python(py, table.columns()[idx].clone());
                    }
                }
                Err(PyKeyError::new_err(name.to_string()))
            }
            _ => unreachable!("Table instances always hold q tables"),
        }
    }

    #[pyo3(signature = (requested_schema=None))]
    fn __arrow_c_stream__(
        slf: PyRef<'_, Self>,
        py: Python<'_>,
        requested_schema: Option<Bound<'_, PyAny>>,
    ) -> PyResult<Py<PyAny>> {
        let schema_capsule: Option<Bound<'_, PyCapsule>> = requested_schema
            .map(|s| s.downcast_into::<PyCapsule>())
            .transpose()?;
        let table = match slf.as_super().inner() {
            qroissant_core::Value::Table(t) => t.clone(),
            _ => unreachable!("Table instances always hold q tables"),
        };
        let opts = effective_export_opts(py, slf.as_super().projection_opts());
        let export = py
            .detach(|| project_table(&table, &opts).map_err(|e| e.to_string()))
            .map_err(|e| PyNotImplementedError::new_err(e))?;
        let reader = ArrayIterator::new(vec![Ok(export.struct_array)], export.struct_field);
        let capsule = to_stream_pycapsule(py, Box::new(reader), schema_capsule)?;
        Ok(capsule.into_any().unbind())
    }

    fn __repr__(slf: PyRef<'_, Self>) -> String {
        match slf.as_super().inner() {
            CoreValue::Table(table) => crate::repr::format_table(table),
            _ => unreachable!("Table instances always hold q tables"),
        }
    }

    fn __str__(slf: PyRef<'_, Self>) -> String {
        Self::__repr__(slf)
    }
}

fn normalize_index(index: isize, len: usize) -> PyResult<usize> {
    let len = len as isize;
    let index = if index < 0 { len + index } else { index };
    if !(0..len).contains(&index) {
        return Err(PyIndexError::new_err("index out of range"));
    }
    Ok(index as usize)
}

fn bytes_or_utf8(value: &Bound<'_, PyAny>) -> PyResult<Vec<u8>> {
    if let Ok(bytes) = value.extract::<Vec<u8>>() {
        return Ok(bytes);
    }
    Ok(value.extract::<String>()?.into_bytes())
}

fn atom_to_python(py: Python<'_>, atom: &CoreAtom) -> PyResult<Py<PyAny>> {
    match atom {
        CoreAtom::Boolean(value) => Ok(value.into_pyobject(py)?.to_owned().unbind().into_any()),
        CoreAtom::Guid(value) => Ok(PyBytes::new(py, value).unbind().into_any()),
        CoreAtom::Byte(value) => Ok(value.into_pyobject(py)?.unbind().into_any()),
        CoreAtom::Short(value) => Ok(value.into_pyobject(py)?.unbind().into_any()),
        CoreAtom::Int(value) => Ok(value.into_pyobject(py)?.unbind().into_any()),
        CoreAtom::Long(value) => Ok(value.into_pyobject(py)?.unbind().into_any()),
        CoreAtom::Real(value) => Ok(value.into_pyobject(py)?.unbind().into_any()),
        CoreAtom::Float(value) => Ok(value.into_pyobject(py)?.unbind().into_any()),
        CoreAtom::Char(value) => Ok(value.into_pyobject(py)?.unbind().into_any()),
        CoreAtom::Symbol(value) => Ok(PyBytes::new(py, value).unbind().into_any()),
        CoreAtom::Timestamp(value) => Ok(value.into_pyobject(py)?.unbind().into_any()),
        CoreAtom::Month(value) => Ok(value.into_pyobject(py)?.unbind().into_any()),
        CoreAtom::Date(value) => Ok(value.into_pyobject(py)?.unbind().into_any()),
        CoreAtom::Datetime(value) => Ok(value.into_pyobject(py)?.unbind().into_any()),
        CoreAtom::Timespan(value) => Ok(value.into_pyobject(py)?.unbind().into_any()),
        CoreAtom::Minute(value) => Ok(value.into_pyobject(py)?.unbind().into_any()),
        CoreAtom::Second(value) => Ok(value.into_pyobject(py)?.unbind().into_any()),
        CoreAtom::Time(value) => Ok(value.into_pyobject(py)?.unbind().into_any()),
    }
}

fn atom_from_python(qtype: &Type, value: &Bound<'_, PyAny>) -> PyResult<CoreAtom> {
    ensure_shape(qtype, Shape::Atom)?;
    let primitive = qtype
        .primitive_value()
        .ok_or_else(|| PyValueError::new_err("atom qtype requires a primitive"))?;
    match primitive {
        Primitive::Boolean => Ok(CoreAtom::Boolean(value.extract()?)),
        Primitive::Guid => {
            let bytes = value.extract::<Vec<u8>>()?;
            let guid: [u8; 16] = bytes.try_into().map_err(|_| {
                PyValueError::new_err("guid atoms must be backed by exactly 16 bytes")
            })?;
            Ok(CoreAtom::Guid(guid))
        }
        Primitive::Byte => Ok(CoreAtom::Byte(value.extract()?)),
        Primitive::Short => Ok(CoreAtom::Short(value.extract()?)),
        Primitive::Int => Ok(CoreAtom::Int(value.extract()?)),
        Primitive::Long => Ok(CoreAtom::Long(value.extract()?)),
        Primitive::Real => Ok(CoreAtom::Real(value.extract()?)),
        Primitive::Float => Ok(CoreAtom::Float(value.extract()?)),
        Primitive::Char => Ok(CoreAtom::Char(extract_char_like(value)?)),
        Primitive::Symbol => Ok(CoreAtom::Symbol(bytes::Bytes::from(bytes_or_utf8(value)?))),
        Primitive::Timestamp => Ok(CoreAtom::Timestamp(value.extract()?)),
        Primitive::Month => Ok(CoreAtom::Month(value.extract()?)),
        Primitive::Date => Ok(CoreAtom::Date(value.extract()?)),
        Primitive::Datetime => Ok(CoreAtom::Datetime(value.extract()?)),
        Primitive::Timespan => Ok(CoreAtom::Timespan(value.extract()?)),
        Primitive::Minute => Ok(CoreAtom::Minute(value.extract()?)),
        Primitive::Second => Ok(CoreAtom::Second(value.extract()?)),
        Primitive::Time => Ok(CoreAtom::Time(value.extract()?)),
        Primitive::Mixed => Err(PyValueError::new_err("mixed atoms are not valid q values")),
    }
}

fn extract_char_like(value: &Bound<'_, PyAny>) -> PyResult<u8> {
    if let Ok(byte) = value.extract::<u8>() {
        return Ok(byte);
    }
    let bytes = value.extract::<Vec<u8>>()?;
    let [byte] = <[u8; 1]>::try_from(bytes.as_slice())
        .map_err(|_| PyValueError::new_err("char values must be a single byte or integer"))?;
    Ok(byte)
}

fn vector_from_python(qtype: &Type, values: &Bound<'_, PyAny>) -> PyResult<CoreVector> {
    ensure_shape(qtype, Shape::Vector)?;
    let primitive = qtype
        .primitive_value()
        .ok_or_else(|| PyValueError::new_err("vector qtype requires a primitive"))?;
    let list = values
        .cast::<PyList>()
        .map_err(|_| PyValueError::new_err("vector payloads must be Python lists"))?;
    let attribute = qtype.attribute_value().unwrap_or(Attribute::None).into();
    let data = match primitive {
        Primitive::Boolean => {
            let bools: Vec<bool> = extract_list(list, |item| item.extract())?;
            let bytes: Vec<u8> = bools.into_iter().map(|b| if b { 1 } else { 0 }).collect();
            VectorData::Boolean(bytes::Bytes::from(bytes))
        }
        Primitive::Guid => VectorData::from_guids(&extract_list(list, |item| {
            let bytes = item.extract::<Vec<u8>>()?;
            bytes
                .try_into()
                .map_err(|_| PyValueError::new_err("guid vector elements must be exactly 16 bytes"))
        })?),
        Primitive::Byte => VectorData::Byte(bytes::Bytes::from(extract_list(list, |item| {
            item.extract::<u8>()
        })?)),
        Primitive::Short => VectorData::from_i16s(&extract_list(list, |item| item.extract())?),
        Primitive::Int => VectorData::from_i32s(&extract_list(list, |item| item.extract())?),
        Primitive::Long => VectorData::from_i64s(&extract_list(list, |item| item.extract())?),
        Primitive::Real => VectorData::from_f32s(&extract_list(list, |item| item.extract())?),
        Primitive::Float => VectorData::from_f64s(&extract_list(list, |item| item.extract())?),
        Primitive::Char => {
            VectorData::Char(bytes::Bytes::from(extract_list(list, extract_char_like)?))
        }
        Primitive::Symbol => VectorData::Symbol(
            extract_list(list, bytes_or_utf8)?
                .into_iter()
                .map(bytes::Bytes::from)
                .collect(),
        ),
        Primitive::Timestamp => {
            VectorData::from_timestamps(&extract_list(list, |item| item.extract())?)
        }
        Primitive::Month => VectorData::from_months(&extract_list(list, |item| item.extract())?),
        Primitive::Date => VectorData::from_dates(&extract_list(list, |item| item.extract())?),
        Primitive::Datetime => {
            VectorData::from_datetimes(&extract_list(list, |item| item.extract())?)
        }
        Primitive::Timespan => {
            VectorData::from_timespans(&extract_list(list, |item| item.extract())?)
        }
        Primitive::Minute => VectorData::from_minutes(&extract_list(list, |item| item.extract())?),
        Primitive::Second => VectorData::from_seconds(&extract_list(list, |item| item.extract())?),
        Primitive::Time => VectorData::from_times(&extract_list(list, |item| item.extract())?),
        Primitive::Mixed => {
            return Err(PyValueError::new_err(
                "mixed vectors must use List rather than Vector",
            ));
        }
    };
    Ok(CoreVector::new(attribute, data))
}

fn list_from_python(qtype: &Type, values: &Bound<'_, PyAny>) -> PyResult<CoreList> {
    ensure_shape(qtype, Shape::List)?;
    let list = values
        .cast::<PyList>()
        .map_err(|_| PyValueError::new_err("list payloads must be Python lists"))?;
    let attribute = qtype.attribute_value().unwrap_or(Attribute::None).into();
    let mut inner = Vec::with_capacity(list.len());
    for item in list.iter() {
        inner.push(python_to_core_value(&item)?);
    }
    Ok(CoreList::new(attribute, inner))
}

fn dictionary_from_python(
    qtype: &Type,
    keys: &Bound<'_, PyAny>,
    values: &Bound<'_, PyAny>,
) -> PyResult<CoreDictionary> {
    ensure_shape(qtype, Shape::Dictionary)?;
    let sorted = qtype.sorted_value().unwrap_or(false);
    let dictionary = CoreDictionary::new(
        sorted,
        python_to_core_value(keys)?,
        python_to_core_value(values)?,
    );
    dictionary
        .validate()
        .map_err(|error| PyValueError::new_err(error.to_string()))?;
    Ok(dictionary)
}

fn table_from_python(qtype: &Type, columns: &Bound<'_, PyAny>) -> PyResult<CoreTable> {
    ensure_shape(qtype, Shape::Table)?;
    let columns = columns
        .cast::<PyDict>()
        .map_err(|_| PyValueError::new_err("table payloads must be Python dicts"))?;
    let attribute = qtype.attribute_value().unwrap_or(Attribute::None).into();
    let mut names = Vec::with_capacity(columns.len());
    let mut values = Vec::with_capacity(columns.len());
    for (name, column) in columns.iter() {
        names.push(bytes::Bytes::from(name.extract::<String>()?.into_bytes()));
        values.push(python_to_core_value(&column)?);
    }
    let table = CoreTable::new(attribute, names, values);
    table
        .validate()
        .map_err(|error| PyValueError::new_err(error.to_string()))?;
    Ok(table)
}

fn ensure_shape(qtype: &Type, expected: Shape) -> PyResult<()> {
    if qtype.shape_value() != expected {
        return Err(PyValueError::new_err(format!(
            "qtype shape {:?} does not match {:?}",
            qtype.shape_value(),
            expected
        )));
    }
    Ok(())
}

fn extract_list<T, F>(items: &Bound<'_, PyList>, convert: F) -> PyResult<Vec<T>>
where
    F: Fn(&Bound<'_, PyAny>) -> PyResult<T>,
{
    let mut values = Vec::with_capacity(items.len());
    for item in items.iter() {
        values.push(convert(&item)?);
    }
    Ok(values)
}

fn vector_to_pylist(py: Python<'_>, vector: &CoreVector) -> PyResult<Py<PyList>> {
    let len = vector.len();
    let mut values = Vec::with_capacity(len);
    for index in 0..len {
        values.push(vector_item_to_python(py, vector, index)?);
    }
    Ok(PyList::new(py, values)?.unbind())
}

fn vector_item_to_python(py: Python<'_>, vector: &CoreVector, index: usize) -> PyResult<Py<PyAny>> {
    let data = vector.data();
    match data {
        VectorData::Boolean(values) => Ok((values[index] != 0)
            .into_pyobject(py)?
            .to_owned()
            .unbind()
            .into_any()),
        VectorData::Guid(values) => {
            let chunk = &values[index * 16..(index + 1) * 16];
            Ok(PyBytes::new(py, chunk).unbind().into_any())
        }
        VectorData::Byte(values) => Ok(values[index].into_pyobject(py)?.unbind().into_any()),
        VectorData::Short(_) => Ok(data.as_i16_slice()[index]
            .into_pyobject(py)?
            .unbind()
            .into_any()),
        VectorData::Int(_)
        | VectorData::Month(_)
        | VectorData::Date(_)
        | VectorData::Minute(_)
        | VectorData::Second(_)
        | VectorData::Time(_) => Ok(data.as_i32_slice()[index]
            .into_pyobject(py)?
            .unbind()
            .into_any()),
        VectorData::Long(_) | VectorData::Timestamp(_) | VectorData::Timespan(_) => Ok(data
            .as_i64_slice()[index]
            .into_pyobject(py)?
            .unbind()
            .into_any()),
        VectorData::Real(_) => Ok(data.as_f32_slice()[index]
            .into_pyobject(py)?
            .unbind()
            .into_any()),
        VectorData::Float(_) | VectorData::Datetime(_) => Ok(data.as_f64_slice()[index]
            .into_pyobject(py)?
            .unbind()
            .into_any()),
        VectorData::Char(values) => Ok(values[index].into_pyobject(py)?.unbind().into_any()),
        VectorData::Symbol(values) => Ok(PyBytes::new(py, &values[index]).unbind().into_any()),
    }
}

fn map_ingestion_error(e: IngestionError) -> PyErr {
    PyValueError::new_err(e.to_string())
}

/// If Polars is loaded in this Python process and the caller has not opted out
/// of the default heterogeneous-list mode (still `Union`), transparently switch
/// to `CoalesceTemporals` so mixed-precision temporal columns surface as a flat
/// `Timestamp(ns)` instead of an Arrow `Union` (which Polars rejects with a
/// hard panic).
///
/// Heuristic: presence of `polars` in `sys.modules` strongly correlates with a
/// Polars consumer for the Arrow PyCapsule export. Users who want faithful
/// `Union` despite Polars being imported should set the mode explicitly via
/// `DecodeOptions.with_coalesce_temporals(False)` — left as a documented
/// limitation, since explicit `Union` and the default are indistinguishable.
fn effective_export_opts(py: Python<'_>, opts: &Arc<ProjectionOptions>) -> Arc<ProjectionOptions> {
    if opts.heterogeneous_list_mode != HeterogeneousListMode::Union {
        return opts.clone();
    }
    let polars_loaded = py
        .import("sys")
        .and_then(|sys| sys.getattr("modules"))
        .and_then(|m| m.contains("polars"))
        .unwrap_or(false);
    if !polars_loaded {
        return opts.clone();
    }
    let mut next = (**opts).clone();
    next.heterogeneous_list_mode = HeterogeneousListMode::CoalesceTemporals;
    Arc::new(next)
}

pub fn python_to_core_value(value: &Bound<'_, PyAny>) -> PyResult<CoreValue> {
    // Try qroissant Value first (it also implements Arrow protocols, so must come first).
    if let Ok(q_value) = value.extract::<PyRef<'_, Value>>() {
        return Ok(q_value.inner().clone());
    }

    // Check Arrow stream protocol (record batches → table).
    if value.hasattr("__arrow_c_stream__")? {
        let capsule_obj = value.getattr("__arrow_c_stream__")?.call0()?;
        let stream_capsule = capsule_obj.downcast::<PyCapsule>().map_err(PyErr::from)?;
        let reader =
            pyo3_arrow::PyRecordBatchReader::from_arrow_pycapsule(stream_capsule)?.into_reader()?;
        let schema = reader.schema();
        let value = ingest_record_batch_reader(schema, reader).map_err(map_ingestion_error)?;
        return Ok(value);
    }

    // Check Arrow array protocol (single array or record batch).
    if value.hasattr("__arrow_c_array__")? {
        // Try extracting as a record batch first.
        if let Ok(record_batch) = value.extract::<pyo3_arrow::PyRecordBatch>() {
            let batch = record_batch.into_inner();
            let value = ingest_record_batch(batch).map_err(map_ingestion_error)?;
            return Ok(value);
        }
        // Fall back to plain array.
        let array: pyo3_arrow::PyArray = value.extract()?;
        let (array, field) = array.into_inner();
        let value = ingest_array(array, field.as_ref()).map_err(map_ingestion_error)?;
        return Ok(value);
    }

    Err(PyNotImplementedError::new_err(
        "encoding non-qroissant values is not implemented yet; \
         pass a qroissant Value or an object implementing the Arrow protocol",
    ))
}

pub fn core_value_to_python(py: Python<'_>, value: CoreValue) -> PyResult<Py<PyAny>> {
    core_value_to_python_with_opts(py, value, Arc::new(ProjectionOptions::default()))
}

pub fn core_value_to_python_with_opts(
    py: Python<'_>,
    value: CoreValue,
    opts: Arc<ProjectionOptions>,
) -> PyResult<Py<PyAny>> {
    match value {
        CoreValue::Atom(atom) => Ok(Py::new(
            py,
            (Atom, Value::new_with_opts(CoreValue::Atom(atom), opts)),
        )?
        .into_any()),
        CoreValue::Vector(vector) => Ok(Py::new(
            py,
            (
                Vector,
                Value::new_with_opts(CoreValue::Vector(vector), opts),
            ),
        )?
        .into_any()),
        CoreValue::List(list) => Ok(Py::new(
            py,
            (List, Value::new_with_opts(CoreValue::List(list), opts)),
        )?
        .into_any()),
        CoreValue::Dictionary(dictionary) => Ok(Py::new(
            py,
            (
                Dictionary,
                Value::new_with_opts(CoreValue::Dictionary(dictionary), opts),
            ),
        )?
        .into_any()),
        CoreValue::Table(table) => Ok(Py::new(
            py,
            (Table, Value::new_with_opts(CoreValue::Table(table), opts)),
        )?
        .into_any()),
        CoreValue::UnaryPrimitive { opcode } => {
            Ok(Py::new(py, Value::new(CoreValue::UnaryPrimitive { opcode }))?.into_any())
        }
    }
}

pub fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
    module.add_class::<Value>()?;
    module.add_class::<Atom>()?;
    module.add_class::<Vector>()?;
    module.add_class::<List>()?;
    module.add_class::<Dictionary>()?;
    module.add_class::<Table>()?;
    Ok(())
}
