//! Arrow projection: converts decoded q `Value` trees into Arrow arrays.

use std::ptr::NonNull;
use std::sync::Arc;

use arrow_array::ArrayRef;
use arrow_array::BinaryArray;
use arrow_array::BooleanArray;
use arrow_buffer::BooleanBuffer;
use arrow_array::Date32Array;
use arrow_array::DurationNanosecondArray;
use arrow_array::FixedSizeBinaryArray;
use arrow_array::Float32Array;
use arrow_array::Float64Array;
use arrow_array::Int16Array;
use arrow_array::Int32Array;
use arrow_array::Int64Array;
use arrow_array::LargeListArray;
use arrow_array::LargeStringArray;
use arrow_array::ListArray;
use arrow_array::RecordBatch;
use arrow_array::StringArray;
use arrow_array::StringViewArray;
use arrow_array::StructArray;
use arrow_array::Time32MillisecondArray;
use arrow_array::Time32SecondArray;
use arrow_array::TimestampMillisecondArray;
use arrow_array::TimestampNanosecondArray;
use arrow_array::UInt8Array;
use arrow_array::UnionArray;
use arrow_array::builder::StringDictionaryBuilder;
use arrow_array::types::Int32Type;
use arrow_buffer::Buffer;
use arrow_buffer::NullBuffer;
use arrow_buffer::OffsetBuffer;
use arrow_buffer::ScalarBuffer;
use arrow_schema::DataType;
use arrow_schema::Field;
use arrow_schema::FieldRef;
use arrow_schema::Fields;
use arrow_schema::Schema;
use arrow_schema::SchemaRef;
use arrow_schema::TimeUnit;
use arrow_schema::UnionFields;
use arrow_schema::UnionMode as ArrowUnionMode;
use chrono::Months;
use chrono::NaiveDate;
use qroissant_core::Atom;
use qroissant_core::Attribute;
use qroissant_core::Dictionary;
use qroissant_core::List;
use qroissant_core::Table;
use qroissant_core::Value;
use qroissant_core::Vector;
use qroissant_core::VectorData;
use qroissant_kernels::nulls::Q_INF_DATE;
use qroissant_kernels::nulls::Q_INF_INT;
use qroissant_kernels::nulls::Q_INF_LONG;
use qroissant_kernels::nulls::Q_INF_MINUTE;
use qroissant_kernels::nulls::Q_INF_SECOND;
use qroissant_kernels::nulls::Q_INF_SHORT;
use qroissant_kernels::nulls::Q_INF_TIME;
use qroissant_kernels::nulls::Q_INF_TIMESPAN;
use qroissant_kernels::nulls::Q_INF_TIMESTAMP;
use qroissant_kernels::nulls::Q_NINF_DATE;
use qroissant_kernels::nulls::Q_NINF_INT;
use qroissant_kernels::nulls::Q_NINF_LONG;
use qroissant_kernels::nulls::Q_NINF_MINUTE;
use qroissant_kernels::nulls::Q_NINF_SECOND;
use qroissant_kernels::nulls::Q_NINF_SHORT;
use qroissant_kernels::nulls::Q_NINF_TIME;
use qroissant_kernels::nulls::Q_NINF_TIMESPAN;
use qroissant_kernels::nulls::Q_NINF_TIMESTAMP;
use qroissant_kernels::nulls::Q_NULL_DATE;
use qroissant_kernels::nulls::Q_NULL_INT;
use qroissant_kernels::nulls::Q_NULL_LONG;
use qroissant_kernels::nulls::Q_NULL_MINUTE;
use qroissant_kernels::nulls::Q_NULL_MONTH;
use qroissant_kernels::nulls::Q_NULL_SECOND;
use qroissant_kernels::nulls::Q_NULL_SHORT;
use qroissant_kernels::nulls::Q_NULL_TIME;
use qroissant_kernels::nulls::Q_NULL_TIMESPAN;
use qroissant_kernels::nulls::Q_NULL_TIMESTAMP;
use qroissant_kernels::nulls::validity_f32;
use qroissant_kernels::nulls::validity_f64;
use qroissant_kernels::nulls::validity_i16;
use qroissant_kernels::nulls::validity_i32;
use qroissant_kernels::nulls::validity_i64;
use qroissant_kernels::temporal::DATE_OFFSET_DAYS;
use qroissant_kernels::temporal::MILLIS_PER_DAY;
use qroissant_kernels::temporal::TIMESTAMP_OFFSET_NS;
use qroissant_kernels::boolean::pack_bool_bytes;
use qroissant_kernels::temporal::copy_and_minutes_to_seconds;
use qroissant_kernels::temporal::copy_and_offset_dates;
use qroissant_kernels::temporal::copy_and_offset_timestamps;
use rayon::prelude::*;

use crate::error::ProjectionError;
use crate::error::ProjectionResult;
use crate::metadata::q_field;
use crate::options::HeterogeneousListMode;
use crate::options::ListProjection;
use crate::options::ProjectionOptions;
use crate::options::StringProjection;
use crate::options::SymbolProjection;
use crate::options::UnionMode;

const MILLENNIUM: NaiveDate = NaiveDate::from_ymd_opt(2000, 1, 1).unwrap();
const UNIX_EPOCH: NaiveDate = NaiveDate::from_ymd_opt(1970, 1, 1).unwrap();

// ---------------------------------------------------------------------------
// Public output types
// ---------------------------------------------------------------------------

/// Projected Arrow array together with its annotated field descriptor.
pub struct ArrayExport {
    pub array: ArrayRef,
    pub field: FieldRef,
}

/// Projected Arrow record batch together with the struct-typed wrapper needed
/// for the C stream interface.
pub struct BatchExport {
    /// The schema of the record batch.
    pub schema: SchemaRef,
    /// The data as an owned RecordBatch.
    pub batch: RecordBatch,
    /// The StructArray version of `batch`, ready to feed into an Arrow stream.
    pub struct_array: ArrayRef,
    /// A struct field that wraps the schema so pyo3-arrow can create a stream
    /// capsule without needing to inspect Arrow internals.
    pub struct_field: FieldRef,
}

// ---------------------------------------------------------------------------
// Entry points
// ---------------------------------------------------------------------------

/// Project any `Value` (except Table) into an `ArrayExport`.
///
/// For atoms the result is a length-1 array.  For vectors the result is a
/// full-length array.  Tables must be projected with [`project_table`].
pub fn project(value: &Value, opts: &ProjectionOptions) -> ProjectionResult<ArrayExport> {
    match value {
        Value::Atom(atom) => project_atom(atom, opts),
        Value::Vector(vector) => project_vector(vector, opts),
        Value::List(list) => project_list(list, opts),
        Value::Dictionary(dict) => project_dictionary(dict, opts),
        Value::Table(table) => {
            let batch = project_table(table, opts)?;
            Ok(ArrayExport {
                array: batch.struct_array,
                field: batch.struct_field,
            })
        }
        Value::UnaryPrimitive { .. } => {
            let array: ArrayRef = Arc::new(arrow_array::NullArray::new(1));
            let field = q_field(DataType::Null, true, "unary_primitive", None, None, None);
            Ok(ArrayExport { array, field })
        }
    }
}

/// Project a q `Table` into a `BatchExport` suitable for the Arrow C stream
/// interface.
pub fn project_table(table: &Table, opts: &ProjectionOptions) -> ProjectionResult<BatchExport> {
    const PARALLEL_THRESHOLD: usize = 4;
    let use_parallel = opts.parallel && table.num_columns() >= PARALLEL_THRESHOLD;

    let column_pairs: Vec<(&[u8], &qroissant_core::Value)> = table
        .column_names()
        .iter()
        .map(|n| n.as_ref())
        .zip(table.columns())
        .collect();

    let results: Vec<ProjectionResult<(FieldRef, ArrayRef)>> = if use_parallel {
        column_pairs
            .par_iter()
            .map(|(name_bytes, column)| {
                let name = String::from_utf8_lossy(name_bytes);
                let export = project(column, opts)?;
                let named_field = Arc::new(export.field.as_ref().clone().with_name(name.as_ref()));
                Ok((named_field, export.array))
            })
            .collect()
    } else {
        column_pairs
            .iter()
            .map(|(name_bytes, column)| {
                let name = String::from_utf8_lossy(name_bytes);
                let export = project(column, opts)?;
                let named_field = Arc::new(export.field.as_ref().clone().with_name(name.as_ref()));
                Ok((named_field, export.array))
            })
            .collect()
    };

    let mut fields: Vec<FieldRef> = Vec::with_capacity(table.num_columns());
    let mut arrays: Vec<ArrayRef> = Vec::with_capacity(table.num_columns());
    for result in results {
        let (field, array) = result?;
        fields.push(field);
        arrays.push(array);
    }

    // Build schema metadata marking this as a q table.
    let mut meta = std::collections::HashMap::new();
    meta.insert(crate::metadata::SHAPE_KEY.to_string(), "table".to_string());
    if table.attribute() != Attribute::None {
        meta.insert(
            crate::metadata::ATTRIBUTE_KEY.to_string(),
            attribute_meta_str(table.attribute()).to_string(),
        );
    }
    let schema = Arc::new(Schema::new(fields.clone()).with_metadata(meta));

    let batch = RecordBatch::try_new(schema.clone(), arrays)
        .map_err(|e| ProjectionError::Arrow(e.to_string()))?;

    // Wrap as StructArray for the C stream interface.
    let struct_field: FieldRef = Arc::new(
        Field::new_struct("", schema.fields().clone(), false)
            .with_metadata(schema.metadata().clone()),
    );
    let struct_array: ArrayRef = Arc::new(arrow_array::StructArray::from(batch.clone()));

    Ok(BatchExport {
        schema,
        batch,
        struct_array,
        struct_field,
    })
}

// ---------------------------------------------------------------------------
// Atom projection
// ---------------------------------------------------------------------------

fn project_atom(atom: &Atom, opts: &ProjectionOptions) -> ProjectionResult<ArrayExport> {
    let primitive = atom.primitive();
    let (array, dt) = match atom {
        Atom::Boolean(v) => {
            let arr: ArrayRef = Arc::new(BooleanArray::from(vec![*v]));
            (arr, DataType::Boolean)
        }
        Atom::Guid(v) => {
            let arr: ArrayRef = Arc::new(
                FixedSizeBinaryArray::try_from_iter(std::iter::once(v.as_slice()))
                    .map_err(|e| ProjectionError::Arrow(e.to_string()))?,
            );
            (arr, DataType::FixedSizeBinary(16))
        }
        Atom::Byte(v) => {
            let arr: ArrayRef = Arc::new(UInt8Array::from(vec![*v]));
            (arr, DataType::UInt8)
        }
        Atom::Short(v) => {
            let opt = if *v == Q_NULL_SHORT { None } else { Some(*v) };
            let arr: ArrayRef = Arc::new(Int16Array::from(vec![opt]));
            (arr, DataType::Int16)
        }
        Atom::Int(v) => {
            let opt = if *v == Q_NULL_INT { None } else { Some(*v) };
            let arr: ArrayRef = Arc::new(Int32Array::from(vec![opt]));
            (arr, DataType::Int32)
        }
        Atom::Long(v) => {
            let opt = if *v == Q_NULL_LONG { None } else { Some(*v) };
            let arr: ArrayRef = Arc::new(Int64Array::from(vec![opt]));
            (arr, DataType::Int64)
        }
        Atom::Real(v) => {
            let opt = if v.is_nan() { None } else { Some(*v) };
            let arr: ArrayRef = Arc::new(Float32Array::from(vec![opt]));
            (arr, DataType::Float32)
        }
        Atom::Float(v) => {
            let opt = if v.is_nan() { None } else { Some(*v) };
            let arr: ArrayRef = Arc::new(Float64Array::from(vec![opt]));
            (arr, DataType::Float64)
        }
        Atom::Char(v) => {
            let (arr, dt) = project_char_atom(*v, opts.string);
            return Ok(ArrayExport {
                array: arr,
                field: q_field(dt, false, "atom", Some(primitive), None, None),
            });
        }
        Atom::Symbol(v) => {
            let (arr, dt) = project_symbol_bytes(
                std::slice::from_ref(v),
                opts.symbol,
                opts.assume_symbol_utf8,
            )?;
            return Ok(ArrayExport {
                array: arr,
                field: q_field(dt, false, "atom", Some(primitive), None, None),
            });
        }
        Atom::Timestamp(v) => {
            let opt = if *v == Q_NULL_TIMESTAMP {
                None
            } else {
                Some(v.saturating_add(TIMESTAMP_OFFSET_NS))
            };
            let arr: ArrayRef = Arc::new(TimestampNanosecondArray::from(vec![opt]));
            (arr, DataType::Timestamp(TimeUnit::Nanosecond, None))
        }
        Atom::Month(v) => {
            let opt = q_month_to_date32(*v);
            let arr: ArrayRef = Arc::new(Date32Array::from(vec![opt]));
            (arr, DataType::Date32)
        }
        Atom::Date(v) => {
            let opt = if *v == Q_NULL_DATE {
                None
            } else {
                Some(v.saturating_add(DATE_OFFSET_DAYS))
            };
            let arr: ArrayRef = Arc::new(Date32Array::from(vec![opt]));
            (arr, DataType::Date32)
        }
        Atom::Datetime(v) => {
            let opt = if v.is_nan() {
                None
            } else {
                Some(datetime_to_ms(*v))
            };
            let arr: ArrayRef = Arc::new(TimestampMillisecondArray::from(vec![opt]));
            (arr, DataType::Timestamp(TimeUnit::Millisecond, None))
        }
        Atom::Timespan(v) => {
            let opt = if *v == Q_NULL_TIMESPAN {
                None
            } else {
                Some(*v)
            };
            let arr: ArrayRef = Arc::new(DurationNanosecondArray::from(vec![opt]));
            (arr, DataType::Duration(TimeUnit::Nanosecond))
        }
        Atom::Minute(v) => {
            let opt = if *v == Q_NULL_MINUTE {
                None
            } else {
                Some(v.saturating_mul(60))
            };
            let arr: ArrayRef = Arc::new(Time32SecondArray::from(vec![opt]));
            (arr, DataType::Time32(TimeUnit::Second))
        }
        Atom::Second(v) => {
            let opt = if *v == Q_NULL_SECOND { None } else { Some(*v) };
            let arr: ArrayRef = Arc::new(Time32SecondArray::from(vec![opt]));
            (arr, DataType::Time32(TimeUnit::Second))
        }
        Atom::Time(v) => {
            let opt = if *v == Q_NULL_TIME { None } else { Some(*v) };
            let arr: ArrayRef = Arc::new(Time32MillisecondArray::from(vec![opt]));
            (arr, DataType::Time32(TimeUnit::Millisecond))
        }
    };

    let nullable = array.null_count() > 0;
    let field = q_field(dt, nullable, "atom", Some(primitive), None, None);
    Ok(ArrayExport { array, field })
}

// ---------------------------------------------------------------------------
// Vector projection
// ---------------------------------------------------------------------------

/// Converts an optional SIMD validity vector into an Arrow `NullBuffer`.
fn to_null_buffer(validity: Option<Vec<bool>>) -> Option<NullBuffer> {
    validity.map(|v| NullBuffer::from(v))
}

/// Creates an Arrow `Buffer` backed by a `bytes::Bytes` without copying.
///
/// The `Bytes` reference count is bumped (via `Arc`) to keep the memory alive
/// as long as Arrow holds the buffer. When Arrow drops the buffer, it drops
/// the `Arc<bytes::Bytes>`, which decrements the refcount.
fn bytes_to_arrow_buffer(bytes: bytes::Bytes) -> Buffer {
    let len = bytes.len();
    if len == 0 {
        return Buffer::from_vec(Vec::<u8>::new());
    }
    let ptr = NonNull::new(bytes.as_ptr() as *mut u8).expect("non-null Bytes pointer");
    // SAFETY: `ptr` is valid for `len` bytes, and the `Arc<bytes::Bytes>` moved
    // into the allocation owner keeps the backing memory alive.
    unsafe { Buffer::from_custom_allocation(ptr, len, Arc::new(bytes)) }
}

/// Merges a null-only validity with additional infinity sentinels.
/// Used when `treat_infinity_as_null` is enabled.
fn merge_infinity_i16(values: &[i16], null_validity: Option<Vec<bool>>) -> Option<Vec<bool>> {
    let has_inf = values
        .iter()
        .any(|&v| v == Q_INF_SHORT || v == Q_NINF_SHORT);
    if !has_inf {
        return null_validity;
    }
    let mut validity = null_validity.unwrap_or_else(|| vec![true; values.len()]);
    for (i, &v) in values.iter().enumerate() {
        if v == Q_INF_SHORT || v == Q_NINF_SHORT {
            validity[i] = false;
        }
    }
    Some(validity)
}

fn merge_infinity_i32(
    values: &[i32],
    sentinel_inf: i32,
    sentinel_ninf: i32,
    null_validity: Option<Vec<bool>>,
) -> Option<Vec<bool>> {
    let has_inf = values
        .iter()
        .any(|&v| v == sentinel_inf || v == sentinel_ninf);
    if !has_inf {
        return null_validity;
    }
    let mut validity = null_validity.unwrap_or_else(|| vec![true; values.len()]);
    for (i, &v) in values.iter().enumerate() {
        if v == sentinel_inf || v == sentinel_ninf {
            validity[i] = false;
        }
    }
    Some(validity)
}

fn merge_infinity_i64(
    values: &[i64],
    sentinel_inf: i64,
    sentinel_ninf: i64,
    null_validity: Option<Vec<bool>>,
) -> Option<Vec<bool>> {
    let has_inf = values
        .iter()
        .any(|&v| v == sentinel_inf || v == sentinel_ninf);
    if !has_inf {
        return null_validity;
    }
    let mut validity = null_validity.unwrap_or_else(|| vec![true; values.len()]);
    for (i, &v) in values.iter().enumerate() {
        if v == sentinel_inf || v == sentinel_ninf {
            validity[i] = false;
        }
    }
    Some(validity)
}

fn merge_infinity_f32(values: &[f32], null_validity: Option<Vec<bool>>) -> Option<Vec<bool>> {
    let has_inf = values.iter().any(|v| v.is_infinite());
    if !has_inf {
        return null_validity;
    }
    let mut validity = null_validity.unwrap_or_else(|| vec![true; values.len()]);
    for (i, v) in values.iter().enumerate() {
        if v.is_infinite() {
            validity[i] = false;
        }
    }
    Some(validity)
}

fn merge_infinity_f64(values: &[f64], null_validity: Option<Vec<bool>>) -> Option<Vec<bool>> {
    let has_inf = values.iter().any(|v| v.is_infinite());
    if !has_inf {
        return null_validity;
    }
    let mut validity = null_validity.unwrap_or_else(|| vec![true; values.len()]);
    for (i, v) in values.iter().enumerate() {
        if v.is_infinite() {
            validity[i] = false;
        }
    }
    Some(validity)
}

fn project_vector(vector: &Vector, opts: &ProjectionOptions) -> ProjectionResult<ArrayExport> {
    let attribute = vector.attribute();
    let primitive = vector.primitive();
    let attr_opt = if attribute == Attribute::None {
        None
    } else {
        Some(attribute)
    };

    let data = vector.data();
    let (array, dt): (ArrayRef, DataType) = match data {
        VectorData::Boolean(values) => {
            let (bitmap, len) = pack_bool_bytes(values.as_ref());
            let buf = Buffer::from_vec(bitmap);
            let arr: ArrayRef =
                Arc::new(BooleanArray::new(BooleanBuffer::new(buf, 0, len), None));
            (arr, DataType::Boolean)
        }
        VectorData::Guid(raw) => {
            let len = raw.len() / 16;
            let buf = bytes_to_arrow_buffer(raw.clone());
            let arr: ArrayRef = Arc::new(FixedSizeBinaryArray::new(16, buf, None));
            debug_assert_eq!(arr.len(), len);
            (arr, DataType::FixedSizeBinary(16))
        }
        VectorData::Byte(raw) => {
            let len = raw.len();
            let buf = bytes_to_arrow_buffer(raw.clone());
            let arr: ArrayRef = Arc::new(UInt8Array::new(ScalarBuffer::new(buf, 0, len), None));
            (arr, DataType::UInt8)
        }
        VectorData::Short(raw) => {
            let values = data.as_i16_slice();
            let mut validity = validity_i16(values);
            if opts.treat_infinity_as_null {
                validity = merge_infinity_i16(values, validity);
            }
            let nulls = to_null_buffer(validity);
            let buf = bytes_to_arrow_buffer(raw.clone());
            let arr: ArrayRef = Arc::new(Int16Array::new(
                ScalarBuffer::new(buf, 0, values.len()),
                nulls,
            ));
            (arr, DataType::Int16)
        }
        VectorData::Int(raw) => {
            let values = data.as_i32_slice();
            let mut validity = validity_i32(values, Q_NULL_INT);
            if opts.treat_infinity_as_null {
                validity = merge_infinity_i32(values, Q_INF_INT, Q_NINF_INT, validity);
            }
            let nulls = to_null_buffer(validity);
            let buf = bytes_to_arrow_buffer(raw.clone());
            let arr: ArrayRef = Arc::new(Int32Array::new(
                ScalarBuffer::new(buf, 0, values.len()),
                nulls,
            ));
            (arr, DataType::Int32)
        }
        VectorData::Long(raw) => {
            let values = data.as_i64_slice();
            let mut validity = validity_i64(values, Q_NULL_LONG);
            if opts.treat_infinity_as_null {
                validity = merge_infinity_i64(values, Q_INF_LONG, Q_NINF_LONG, validity);
            }
            let nulls = to_null_buffer(validity);
            let buf = bytes_to_arrow_buffer(raw.clone());
            let arr: ArrayRef = Arc::new(Int64Array::new(
                ScalarBuffer::new(buf, 0, values.len()),
                nulls,
            ));
            (arr, DataType::Int64)
        }
        VectorData::Real(raw) => {
            let values = data.as_f32_slice();
            let mut validity = validity_f32(values);
            if opts.treat_infinity_as_null {
                validity = merge_infinity_f32(values, validity);
            }
            let nulls = to_null_buffer(validity);
            let buf = bytes_to_arrow_buffer(raw.clone());
            let arr: ArrayRef = Arc::new(Float32Array::new(
                ScalarBuffer::new(buf, 0, values.len()),
                nulls,
            ));
            (arr, DataType::Float32)
        }
        VectorData::Float(raw) => {
            let values = data.as_f64_slice();
            let mut validity = validity_f64(values);
            if opts.treat_infinity_as_null {
                validity = merge_infinity_f64(values, validity);
            }
            let nulls = to_null_buffer(validity);
            let buf = bytes_to_arrow_buffer(raw.clone());
            let arr: ArrayRef = Arc::new(Float64Array::new(
                ScalarBuffer::new(buf, 0, values.len()),
                nulls,
            ));
            (arr, DataType::Float64)
        }
        VectorData::Char(values) => {
            let (arr, dt) = project_char_vector(values, opts.string);
            return Ok(ArrayExport {
                array: arr,
                field: q_field(dt, false, "vector", Some(primitive), attr_opt, None),
            });
        }
        VectorData::Symbol(values) => {
            let (arr, dt) = project_symbol_bytes(values, opts.symbol, opts.assume_symbol_utf8)?;
            return Ok(ArrayExport {
                array: arr,
                field: q_field(dt, false, "vector", Some(primitive), attr_opt, None),
            });
        }
        VectorData::Timestamp(_) => {
            let src = data.as_i64_slice();
            let mut validity = validity_i64(src, Q_NULL_TIMESTAMP);
            if opts.treat_infinity_as_null {
                validity = merge_infinity_i64(src, Q_INF_TIMESTAMP, Q_NINF_TIMESTAMP, validity);
            }
            let nulls = to_null_buffer(validity);
            let mut typed_buf = vec![0i64; src.len()];
            copy_and_offset_timestamps(src, &mut typed_buf);
            let buf = Buffer::from_vec(typed_buf);
            let arr: ArrayRef = Arc::new(TimestampNanosecondArray::new(
                ScalarBuffer::new(buf, 0, src.len()),
                nulls,
            ));
            (arr, DataType::Timestamp(TimeUnit::Nanosecond, None))
        }
        VectorData::Month(_) => {
            let values = data.as_i32_slice();
            // Month conversion is non-linear (variable days per month), keep scalar path.
            let arr: ArrayRef = Arc::new(Date32Array::from_iter(
                values.iter().map(|&v| q_month_to_date32(v)),
            ));
            (arr, DataType::Date32)
        }
        VectorData::Date(_) => {
            let src = data.as_i32_slice();
            let mut validity = validity_i32(src, Q_NULL_DATE);
            if opts.treat_infinity_as_null {
                validity = merge_infinity_i32(src, Q_INF_DATE, Q_NINF_DATE, validity);
            }
            let nulls = to_null_buffer(validity);
            let mut typed_buf = vec![0i32; src.len()];
            copy_and_offset_dates(src, &mut typed_buf);
            let buf = Buffer::from_vec(typed_buf);
            let arr: ArrayRef =
                Arc::new(Date32Array::new(ScalarBuffer::new(buf, 0, src.len()), nulls));
            (arr, DataType::Date32)
        }
        VectorData::Datetime(_) => {
            let values = data.as_f64_slice();
            // Datetime is f64 fractional days — non-trivial transform, keep scalar path.
            let nulls = to_null_buffer(validity_f64(values));
            let ms_values: Vec<i64> = values
                .iter()
                .map(|&v| if v.is_nan() { 0 } else { datetime_to_ms(v) })
                .collect();
            let arr: ArrayRef = Arc::new(TimestampMillisecondArray::new(
                ScalarBuffer::from(ms_values),
                nulls,
            ));
            (arr, DataType::Timestamp(TimeUnit::Millisecond, None))
        }
        VectorData::Timespan(raw) => {
            let values = data.as_i64_slice();
            let mut validity = validity_i64(values, Q_NULL_TIMESPAN);
            if opts.treat_infinity_as_null {
                validity = merge_infinity_i64(values, Q_INF_TIMESPAN, Q_NINF_TIMESPAN, validity);
            }
            let nulls = to_null_buffer(validity);
            let buf = bytes_to_arrow_buffer(raw.clone());
            let arr: ArrayRef = Arc::new(DurationNanosecondArray::new(
                ScalarBuffer::new(buf, 0, values.len()),
                nulls,
            ));
            (arr, DataType::Duration(TimeUnit::Nanosecond))
        }
        VectorData::Minute(_) => {
            let src = data.as_i32_slice();
            let mut validity = validity_i32(src, Q_NULL_MINUTE);
            if opts.treat_infinity_as_null {
                validity = merge_infinity_i32(src, Q_INF_MINUTE, Q_NINF_MINUTE, validity);
            }
            let nulls = to_null_buffer(validity);
            let mut typed_buf = vec![0i32; src.len()];
            copy_and_minutes_to_seconds(src, &mut typed_buf);
            let buf = Buffer::from_vec(typed_buf);
            let arr: ArrayRef = Arc::new(Time32SecondArray::new(
                ScalarBuffer::new(buf, 0, src.len()),
                nulls,
            ));
            (arr, DataType::Time32(TimeUnit::Second))
        }
        VectorData::Second(raw) => {
            let values = data.as_i32_slice();
            let mut validity = validity_i32(values, Q_NULL_SECOND);
            if opts.treat_infinity_as_null {
                validity = merge_infinity_i32(values, Q_INF_SECOND, Q_NINF_SECOND, validity);
            }
            let nulls = to_null_buffer(validity);
            let buf = bytes_to_arrow_buffer(raw.clone());
            let arr: ArrayRef = Arc::new(Time32SecondArray::new(
                ScalarBuffer::new(buf, 0, values.len()),
                nulls,
            ));
            (arr, DataType::Time32(TimeUnit::Second))
        }
        VectorData::Time(raw) => {
            let values = data.as_i32_slice();
            let mut validity = validity_i32(values, Q_NULL_TIME);
            if opts.treat_infinity_as_null {
                validity = merge_infinity_i32(values, Q_INF_TIME, Q_NINF_TIME, validity);
            }
            let nulls = to_null_buffer(validity);
            let buf = bytes_to_arrow_buffer(raw.clone());
            let arr: ArrayRef = Arc::new(Time32MillisecondArray::new(
                ScalarBuffer::new(buf, 0, values.len()),
                nulls,
            ));
            (arr, DataType::Time32(TimeUnit::Millisecond))
        }
    };

    let nullable = array.null_count() > 0;
    let field = q_field(dt, nullable, "vector", Some(primitive), attr_opt, None);
    Ok(ArrayExport { array, field })
}

// ---------------------------------------------------------------------------
// Symbol and char projection helpers
// ---------------------------------------------------------------------------

/// Project a slice of q symbol byte-strings using the requested `SymbolProjection`.
fn project_symbol_bytes(
    values: &[bytes::Bytes],
    mode: SymbolProjection,
    assume_utf8: bool,
) -> ProjectionResult<(ArrayRef, DataType)> {
    match mode {
        SymbolProjection::Utf8 => {
            let arr: ArrayRef = if assume_utf8 {
                Arc::new(StringArray::from_iter_values(values.iter().map(|v| {
                    // SAFETY: caller asserts symbols are valid UTF-8 (q symbols are ASCII).
                    unsafe { std::str::from_utf8_unchecked(v) }
                })))
            } else {
                Arc::new(StringArray::from_iter_values(
                    values
                        .iter()
                        .map(|v| String::from_utf8_lossy(v).into_owned()),
                ))
            };
            Ok((arr, DataType::Utf8))
        }
        SymbolProjection::LargeUtf8 => {
            let arr: ArrayRef = if assume_utf8 {
                Arc::new(LargeStringArray::from_iter_values(
                    values
                        .iter()
                        .map(|v| unsafe { std::str::from_utf8_unchecked(v) }),
                ))
            } else {
                Arc::new(LargeStringArray::from_iter_values(
                    values
                        .iter()
                        .map(|v| String::from_utf8_lossy(v).into_owned()),
                ))
            };
            Ok((arr, DataType::LargeUtf8))
        }
        SymbolProjection::Utf8View => {
            let arr: ArrayRef = if assume_utf8 {
                Arc::new(StringViewArray::from_iter_values(
                    values
                        .iter()
                        .map(|v| unsafe { std::str::from_utf8_unchecked(v) }),
                ))
            } else {
                Arc::new(StringViewArray::from_iter_values(
                    values
                        .iter()
                        .map(|v| String::from_utf8_lossy(v).into_owned()),
                ))
            };
            Ok((arr, DataType::Utf8View))
        }
        SymbolProjection::Dictionary => {
            let mut builder = StringDictionaryBuilder::<Int32Type>::new();
            if assume_utf8 {
                for v in values {
                    builder.append_value(unsafe { std::str::from_utf8_unchecked(v.as_ref()) });
                }
            } else {
                for v in values {
                    builder.append_value(String::from_utf8_lossy(v.as_ref()));
                }
            }
            let arr: ArrayRef = Arc::new(builder.finish());
            Ok((
                arr,
                DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
            ))
        }
        SymbolProjection::RawBytes => {
            let arr: ArrayRef = Arc::new(BinaryArray::from_iter_values(
                values.iter().map(|v| v.as_ref()),
            ));
            Ok((arr, DataType::Binary))
        }
    }
}

/// Project a single q char atom byte.
fn project_char_atom(v: u8, mode: StringProjection) -> (ArrayRef, DataType) {
    match mode {
        StringProjection::Utf8 => {
            let s = std::str::from_utf8(std::slice::from_ref(&v)).unwrap_or("\u{FFFD}");
            let arr: ArrayRef = Arc::new(StringArray::from(vec![s]));
            (arr, DataType::Utf8)
        }
        StringProjection::Binary => {
            let arr: ArrayRef = Arc::new(BinaryArray::from_iter_values(std::iter::once(
                std::slice::from_ref(&v),
            )));
            (arr, DataType::Binary)
        }
    }
}

/// Project a q char vector as a per-character string/binary array.
fn project_char_vector(values: &[u8], mode: StringProjection) -> (ArrayRef, DataType) {
    match mode {
        StringProjection::Utf8 => {
            let arr: ArrayRef = Arc::new(StringArray::from_iter_values(
                values.iter().enumerate().map(|(i, &b)| {
                    if b.is_ascii() {
                        // SAFETY: a single ASCII byte is valid UTF-8.
                        unsafe { std::str::from_utf8_unchecked(&values[i..i + 1]) }
                    } else {
                        "\u{FFFD}"
                    }
                }),
            ));
            (arr, DataType::Utf8)
        }
        StringProjection::Binary => {
            let arr: ArrayRef = Arc::new(BinaryArray::from_iter_values(
                values.iter().map(std::slice::from_ref),
            ));
            (arr, DataType::Binary)
        }
    }
}

// ---------------------------------------------------------------------------
// List projection
// ---------------------------------------------------------------------------

fn project_list(list: &List, opts: &ProjectionOptions) -> ProjectionResult<ArrayExport> {
    let values = list.values();
    if values.is_empty() {
        // Empty list: LargeList<null> or List<null> depending on opts.
        let child_field: FieldRef = Arc::new(Field::new("item", DataType::Null, true));
        return match opts.list {
            ListProjection::List | ListProjection::ListView => {
                let offsets = OffsetBuffer::<i32>::new(ScalarBuffer::from(vec![0i32]));
                let child_array: ArrayRef = Arc::new(arrow_array::NullArray::new(0));
                let array: ArrayRef = Arc::new(ListArray::new(
                    child_field.clone(),
                    offsets,
                    child_array,
                    None,
                ));
                let field = q_field(DataType::List(child_field), false, "list", None, None, None);
                Ok(ArrayExport { array, field })
            }
            ListProjection::LargeList => {
                let offsets = OffsetBuffer::<i64>::new(ScalarBuffer::from(vec![0i64]));
                let child_array: ArrayRef = Arc::new(arrow_array::NullArray::new(0));
                let array: ArrayRef = Arc::new(LargeListArray::new(
                    child_field.clone(),
                    offsets,
                    child_array,
                    None,
                ));
                let field = q_field(
                    DataType::LargeList(child_field),
                    false,
                    "list",
                    None,
                    None,
                    None,
                );
                Ok(ArrayExport { array, field })
            }
        };
    }

    // Project each element; they may be heterogeneous.
    let mut child_exports: Vec<ArrayExport> = Vec::with_capacity(values.len());
    for v in values {
        child_exports.push(project(v, opts)?);
    }

    // Check homogeneity: all child arrays must share the same Arrow DataType.
    let first_dt = child_exports[0].array.data_type().clone();
    let homogeneous = child_exports
        .iter()
        .all(|e| e.array.data_type() == &first_dt);

    if !homogeneous {
        return project_heterogeneous_list(values, child_exports, opts);
    }

    // Concatenate child arrays.
    let refs: Vec<&dyn arrow_array::Array> =
        child_exports.iter().map(|e| e.array.as_ref()).collect();
    let concatenated =
        arrow_select::concat::concat(&refs).map_err(|e| ProjectionError::Arrow(e.to_string()))?;

    let child_field: FieldRef = Arc::new(Field::new("item", first_dt.clone(), true));
    let attr_opt = if list.attribute() == Attribute::None {
        None
    } else {
        Some(list.attribute())
    };

    match opts.list {
        ListProjection::List | ListProjection::ListView => {
            // Build i32 offsets.
            let mut offsets: Vec<i32> = Vec::with_capacity(values.len() + 1);
            offsets.push(0);
            for e in &child_exports {
                offsets.push(*offsets.last().unwrap() + e.array.len() as i32);
            }
            let offset_buf = OffsetBuffer::<i32>::new(ScalarBuffer::from(offsets));
            let array: ArrayRef = Arc::new(ListArray::new(
                child_field.clone(),
                offset_buf,
                concatenated,
                None,
            ));
            let field = q_field(
                DataType::List(child_field),
                false,
                "list",
                None,
                attr_opt,
                None,
            );
            Ok(ArrayExport { array, field })
        }
        ListProjection::LargeList => {
            let mut offsets: Vec<i64> = Vec::with_capacity(values.len() + 1);
            offsets.push(0);
            for e in &child_exports {
                offsets.push(*offsets.last().unwrap() + e.array.len() as i64);
            }
            let offset_buf = OffsetBuffer::<i64>::new(ScalarBuffer::from(offsets));
            let array: ArrayRef = Arc::new(LargeListArray::new(
                child_field.clone(),
                offset_buf,
                concatenated,
                None,
            ));
            let field = q_field(
                DataType::LargeList(child_field),
                false,
                "list",
                None,
                attr_opt,
                None,
            );
            Ok(ArrayExport { array, field })
        }
    }
}

// ---------------------------------------------------------------------------
// Heterogeneous list → Union projection
// ---------------------------------------------------------------------------

fn is_temporal_dt(dt: &DataType) -> bool {
    matches!(
        dt,
        DataType::Timestamp(_, _)
            | DataType::Date32
            | DataType::Date64
            | DataType::Time32(_)
            | DataType::Time64(_)
            | DataType::Duration(_)
    )
}

/// Attempts to coalesce a heterogeneous list of temporal Arrow arrays into a
/// flat `Timestamp(Nanosecond, None)` array. Returns `None` if any arm is
/// non-temporal or `arrow_cast::cast` rejects an arm. Callers fall back to
/// Union encoding in that case.
fn try_coalesce_temporals_to_ns(child_exports: &[ArrayExport]) -> Option<ArrayRef> {
    let target = DataType::Timestamp(TimeUnit::Nanosecond, None);
    let mut cast_arrays: Vec<ArrayRef> = Vec::with_capacity(child_exports.len());
    for export in child_exports {
        let dt = export.array.data_type();
        if !is_temporal_dt(dt) {
            return None;
        }
        match arrow_cast::cast(export.array.as_ref(), &target) {
            Ok(arr) => cast_arrays.push(arr),
            Err(_) => return None,
        }
    }
    let refs: Vec<&dyn arrow_array::Array> = cast_arrays.iter().map(|a| a.as_ref()).collect();
    arrow_select::concat::concat(&refs).ok()
}

fn project_heterogeneous_list(
    values: &[Value],
    child_exports: Vec<ArrayExport>,
    opts: &ProjectionOptions,
) -> ProjectionResult<ArrayExport> {
    // Opt-in: collapse all-temporal heterogeneous lists into Timestamp(ns).
    // Polars and similar consumers cannot ingest Arrow Union; this lets users
    // accept the lossy promotion in exchange for a flat, ingestible column.
    if opts.heterogeneous_list_mode == HeterogeneousListMode::CoalesceTemporals
        && let Some(flat) = try_coalesce_temporals_to_ns(&child_exports)
    {
        let dt = DataType::Timestamp(TimeUnit::Nanosecond, None);
        let field = q_field(dt, false, "list", None, None, None);
        return Ok(ArrayExport { array: flat, field });
    }

    // Assign a type_id (i8) to each unique Arrow DataType in insertion order.
    let mut type_id_map: Vec<(DataType, i8)> = Vec::new();
    let mut type_ids: Vec<i8> = Vec::with_capacity(values.len());

    for export in &child_exports {
        let dt = export.array.data_type().clone();
        let type_id = if let Some(pos) = type_id_map.iter().position(|(d, _)| d == &dt) {
            type_id_map[pos].1
        } else {
            let id = type_id_map.len() as i8;
            type_id_map.push((dt, id));
            id
        };
        type_ids.push(type_id);
    }

    let n_types = type_id_map.len();

    // Build UnionFields from the first occurrence of each type's field descriptor.
    let mut union_fields_vec: Vec<(i8, FieldRef)> = Vec::with_capacity(n_types);
    for (type_id, _) in &type_id_map {
        // Find the first child export with this DataType.
        let first = child_exports
            .iter()
            .find(|e| e.array.data_type() == type_id)
            .expect("type_id_map entry must have matching export");
        union_fields_vec.push((
            type_id_map
                .iter()
                .find(|(d, _)| d == type_id)
                .map(|(_, id)| *id)
                .unwrap(),
            first.field.clone(),
        ));
    }
    let union_fields = UnionFields::try_new(
        union_fields_vec.iter().map(|(id, _)| *id),
        union_fields_vec.iter().map(|(_, f)| f.as_ref().clone()),
    )
    .map_err(|e| ProjectionError::Arrow(e.to_string()))?;

    let arrow_union_mode = match opts.union_mode {
        UnionMode::Dense => ArrowUnionMode::Dense,
        UnionMode::Sparse => ArrowUnionMode::Sparse,
    };

    let union_array = match opts.union_mode {
        UnionMode::Dense => {
            // Dense union: per-type child arrays with per-type offsets.
            let mut per_type_arrays: Vec<Vec<ArrayRef>> = vec![Vec::new(); n_types];
            let mut offsets: Vec<i32> = Vec::with_capacity(values.len());
            for (export, &type_id) in child_exports.iter().zip(type_ids.iter()) {
                let type_idx = type_id as usize;
                offsets.push(per_type_arrays[type_idx].len() as i32);
                per_type_arrays[type_idx].push(export.array.clone());
            }
            // Concatenate per-type.
            let mut children: Vec<ArrayRef> = Vec::with_capacity(n_types);
            for (type_idx, arrays) in per_type_arrays.iter().enumerate() {
                let refs: Vec<&dyn arrow_array::Array> =
                    arrays.iter().map(|a| a.as_ref()).collect();
                let concatenated = if refs.is_empty() {
                    arrow_array::new_empty_array(&type_id_map[type_idx].0)
                } else {
                    arrow_select::concat::concat(&refs)
                        .map_err(|e| ProjectionError::Arrow(e.to_string()))?
                };
                children.push(concatenated);
            }
            UnionArray::try_new(
                union_fields,
                ScalarBuffer::from(type_ids),
                Some(ScalarBuffer::from(offsets)),
                children,
            )
            .map_err(|e| ProjectionError::Arrow(e.to_string()))?
        }
        UnionMode::Sparse => {
            // Sparse union: all children have length = n_elements.
            // For each type's child, fill in the real values at the positions
            // where that type was selected, and nulls elsewhere.
            let n = values.len();
            let mut per_type_builders: Vec<Vec<Option<ArrayRef>>> = vec![vec![None; n]; n_types];
            for (i, (export, &type_id)) in child_exports.iter().zip(type_ids.iter()).enumerate() {
                per_type_builders[type_id as usize][i] = Some(export.array.clone());
            }
            let mut children: Vec<ArrayRef> = Vec::with_capacity(n_types);
            for (type_idx, slots) in per_type_builders.iter().enumerate() {
                let dt = &type_id_map[type_idx].0;
                let mut pieces: Vec<ArrayRef> = Vec::with_capacity(n);
                for slot in slots {
                    if let Some(arr) = slot {
                        pieces.push(arr.clone());
                    } else {
                        // Null slot: single-element null array of the right type.
                        let null = arrow_array::new_null_array(dt, 1);
                        pieces.push(null);
                    }
                }
                let refs: Vec<&dyn arrow_array::Array> =
                    pieces.iter().map(|a| a.as_ref()).collect();
                let concatenated = arrow_select::concat::concat(&refs)
                    .map_err(|e| ProjectionError::Arrow(e.to_string()))?;
                children.push(concatenated);
            }
            UnionArray::try_new(union_fields, ScalarBuffer::from(type_ids), None, children)
                .map_err(|e| ProjectionError::Arrow(e.to_string()))?
        }
    };

    let union_dt = DataType::Union(union_array.fields().clone(), arrow_union_mode);
    let field = q_field(union_dt, false, "list", None, None, None);
    let array: ArrayRef = Arc::new(union_array);
    Ok(ArrayExport { array, field })
}

// ---------------------------------------------------------------------------
// Dictionary projection
// ---------------------------------------------------------------------------

fn project_dictionary(
    dict: &Dictionary,
    opts: &ProjectionOptions,
) -> ProjectionResult<ArrayExport> {
    let keys_export = project(dict.keys(), opts)?;
    let values_export = project(dict.values(), opts)?;

    let keys_field = Arc::new(keys_export.field.as_ref().clone().with_name("keys"));
    let values_field = Arc::new(values_export.field.as_ref().clone().with_name("values"));

    let fields: Fields = Fields::from(vec![keys_field, values_field]);
    let array: ArrayRef = Arc::new(
        StructArray::try_new(
            fields.clone(),
            vec![keys_export.array, values_export.array],
            None,
        )
        .map_err(|e| ProjectionError::Arrow(e.to_string()))?,
    );

    let sorted = dict.sorted();
    let mut meta = std::collections::HashMap::new();
    meta.insert(crate::metadata::SHAPE_KEY.to_string(), "dict".to_string());
    if sorted {
        meta.insert("sorted".to_string(), "true".to_string());
    }
    let struct_field: FieldRef =
        Arc::new(Field::new_struct("dict", fields, false).with_metadata(meta));
    Ok(ArrayExport {
        array,
        field: struct_field,
    })
}

// ---------------------------------------------------------------------------
// Temporal conversion helpers
// ---------------------------------------------------------------------------

/// Convert a q month offset (months since 2000.01) to an Arrow `Date32` value
/// (days since 1970.01.01).  Returns `None` for null months.
fn q_month_to_date32(q_month: i32) -> Option<i32> {
    if q_month == Q_NULL_MONTH {
        return None;
    }
    let date = if q_month >= 0 {
        MILLENNIUM.checked_add_months(Months::new(q_month as u32))?
    } else {
        MILLENNIUM.checked_sub_months(Months::new((-q_month) as u32))?
    };
    // Days since 1970-01-01
    Some((date - UNIX_EPOCH).num_days() as i32)
}

/// Convert a q datetime (fractional days since 2000.01.01) to Arrow
/// `TimestampMillisecond` (milliseconds since 1970.01.01).
#[inline]
fn datetime_to_ms(q_datetime: f64) -> i64 {
    ((q_datetime + DATE_OFFSET_DAYS as f64) * MILLIS_PER_DAY) as i64
}

fn attribute_meta_str(a: Attribute) -> &'static str {
    match a {
        Attribute::None => "none",
        Attribute::Sorted => "sorted",
        Attribute::Unique => "unique",
        Attribute::Parted => "parted",
        Attribute::Grouped => "grouped",
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use qroissant_core::Atom;
    use qroissant_core::List;

    fn mixed_temporal_list() -> Value {
        Value::List(List::new(
            Attribute::None,
            vec![
                Value::Atom(Atom::Timestamp(0)),
                Value::Atom(Atom::Datetime(0.0)),
            ],
        ))
    }

    #[test]
    fn heterogeneous_temporals_emit_union_by_default() {
        let value = mixed_temporal_list();
        let opts = ProjectionOptions::default();
        let export = project(&value, &opts).expect("projection");
        assert!(
            matches!(export.array.data_type(), DataType::Union(_, _)),
            "expected Union under default mode, got {:?}",
            export.array.data_type()
        );
    }

    #[test]
    fn heterogeneous_temporals_coalesce_to_timestamp_ns() {
        let value = mixed_temporal_list();
        let opts = ProjectionOptions {
            heterogeneous_list_mode: HeterogeneousListMode::CoalesceTemporals,
            ..ProjectionOptions::default()
        };
        let export = project(&value, &opts).expect("projection");
        assert_eq!(
            export.array.data_type(),
            &DataType::Timestamp(TimeUnit::Nanosecond, None),
            "expected flat Timestamp(ns) under CoalesceTemporals"
        );
        assert_eq!(export.array.len(), 2, "expected 2 rows, one per arm");
    }

    #[test]
    fn coalesce_falls_back_to_union_for_non_temporal_arm() {
        // Mixed temporal + int: not all-temporal, should fall through to Union.
        let value = Value::List(List::new(
            Attribute::None,
            vec![
                Value::Atom(Atom::Timestamp(0)),
                Value::Atom(Atom::Int(42)),
            ],
        ));
        let opts = ProjectionOptions {
            heterogeneous_list_mode: HeterogeneousListMode::CoalesceTemporals,
            ..ProjectionOptions::default()
        };
        let export = project(&value, &opts).expect("projection");
        assert!(
            matches!(export.array.data_type(), DataType::Union(_, _)),
            "non-temporal arm must force Union fallback, got {:?}",
            export.array.data_type()
        );
    }
}
