ex_arrow

//! IPC stream read/write NIFs, plus record-batch column buffer helpers for Nx.

use std::io::Cursor;
use std::sync::Arc;

use arrow::array::{Int64Array, StringArray};
use arrow::record_batch::RecordBatch;
use arrow_array::types::{
    Float32Type, Float64Type, Int16Type, Int32Type, Int64Type, Int8Type, UInt16Type, UInt32Type,
    UInt64Type, UInt8Type,
};
use arrow_array::{Array, ArrayRef, PrimitiveArray};
use arrow_buffer::Buffer;
use arrow_data::ArrayData;
use arrow_schema::{DataType, Field, Schema};

use arrow_ipc::reader::{FileReader, StreamReader};
use arrow_ipc::writer::{FileWriter, StreamWriter};
use rustler::resource::ResourceArc;
use rustler::{Encoder, Env, Term};

use crate::resources::{ExArrowIpcFile, ExArrowIpcStream, ExArrowRecordBatch, ExArrowSchema, IpcFileBacking, IpcStreamBacking};
use crate::util::{err_encode, ok_encode};

/// Builds a small IPC stream fixture (schema: id int64, name utf8; 2 rows) for tests.
#[rustler::nif]
pub fn ipc_test_fixture_binary<'a>(env: Env<'a>) -> Term<'a> {
    let schema = Arc::new(Schema::new(vec![
        Field::new("id", DataType::Int64, false),
        Field::new("name", DataType::Utf8, false),
    ]));
    let id_array = Arc::new(Int64Array::from(vec![1_i64, 2]));
    let name_array = Arc::new(StringArray::from(vec!["a", "b"]));
    let batch = match RecordBatch::try_new(schema.clone(), vec![id_array, name_array]) {
        Ok(b) => b,
        Err(e) => return err_encode(env, &e.to_string()),
    };
    let mut buf = Vec::new();
    let mut writer = match StreamWriter::try_new(&mut buf, schema.as_ref()) {
        Ok(w) => w,
        Err(e) => return err_encode(env, &e.to_string()),
    };
    if let Err(e) = writer.write(&batch) {
        return err_encode(env, &e.to_string());
    }
    if let Err(e) = writer.finish() {
        return err_encode(env, &e.to_string());
    }
    let mut owned = match rustler::OwnedBinary::new(buf.len()) {
        Some(b) => b,
        None => return err_encode(env, "binary alloc"),
    };
    owned.as_mut_slice().copy_from_slice(&buf);
    let binary = rustler::Binary::from_owned(owned, env);
    ok_encode(env, binary)
}

/// Builds a small IPC file-format fixture (same schema as stream fixture) for tests.
#[rustler::nif]
pub fn ipc_test_fixture_file_binary<'a>(env: Env<'a>) -> Term<'a> {
    let schema = Arc::new(Schema::new(vec![
        Field::new("id", DataType::Int64, false),
        Field::new("name", DataType::Utf8, false),
    ]));
    let id_array = Arc::new(Int64Array::from(vec![1_i64, 2]));
    let name_array = Arc::new(StringArray::from(vec!["a", "b"]));
    let batch = match RecordBatch::try_new(schema.clone(), vec![id_array, name_array]) {
        Ok(b) => b,
        Err(e) => return err_encode(env, &e.to_string()),
    };
    let buf = Vec::new();
    let cursor = Cursor::new(buf);
    let mut writer = match FileWriter::try_new(cursor, schema.as_ref()) {
        Ok(w) => w,
        Err(e) => return err_encode(env, &e.to_string()),
    };
    if let Err(e) = writer.write(&batch) {
        return err_encode(env, &e.to_string());
    }
    let cursor = match writer.into_inner() {
        Ok(c) => c,
        Err(e) => return err_encode(env, &e.to_string()),
    };
    let buf = cursor.into_inner();
    let mut owned = match rustler::OwnedBinary::new(buf.len()) {
        Some(b) => b,
        None => return err_encode(env, "binary alloc"),
    };
    owned.as_mut_slice().copy_from_slice(&buf);
    let binary = rustler::Binary::from_owned(owned, env);
    ok_encode(env, binary)
}

/// Read IPC stream from binary. Returns {:ok, stream_ref} or {:error, msg}.
#[rustler::nif]
pub fn ipc_reader_from_binary<'a>(env: Env<'a>, data: rustler::Binary) -> Term<'a> {
    let bytes = data.as_slice().to_vec();
    let cursor = Cursor::new(bytes);
    match StreamReader::try_new(cursor, None) {
        Ok(reader) => {
            let stream = ExArrowIpcStream {
                reader: IpcStreamBacking::Binary(std::sync::Mutex::new(reader)),
            };
            ok_encode(env, ResourceArc::new(stream))
        }
        Err(e) => err_encode(env, &e.to_string()),
    }
}

/// Read IPC stream from file path (streaming: does not load entire file into memory).
/// Returns {:ok, stream_ref} or {:error, msg}.
#[rustler::nif]
pub fn ipc_reader_from_file<'a>(env: Env<'a>, path: String) -> Term<'a> {
    let file = match std::fs::File::open(&path) {
        Ok(f) => f,
        Err(e) => return err_encode(env, &e.to_string()),
    };
    let reader = match StreamReader::try_new_buffered(file, None) {
        Ok(r) => r,
        Err(e) => return err_encode(env, &e.to_string()),
    };
    let stream = ExArrowIpcStream {
        reader: IpcStreamBacking::File(std::sync::Mutex::new(reader)),
    };
    ok_encode(env, ResourceArc::new(stream))
}

/// Return list of fields for a schema: [{name, type_atom}, ...]. type_atom is :int64, :float64, :utf8, :binary, :boolean, :null, etc.
#[rustler::nif]
pub fn schema_fields<'a>(env: Env<'a>, schema: ResourceArc<ExArrowSchema>) -> Term<'a> {
    let fields: Vec<Term> = schema
        .schema
        .fields()
        .iter()
        .map(|f| {
            (
                f.name().encode(env),
                data_type_to_atom(env, f.data_type()).encode(env),
            )
                .encode(env)
        })
        .collect();
    fields.encode(env)
}

fn data_type_to_atom(env: Env, dt: &arrow_schema::DataType) -> rustler::types::atom::Atom {
    let s = match dt {
        arrow_schema::DataType::Null => "null",
        arrow_schema::DataType::Boolean => "boolean",
        arrow_schema::DataType::Int64 => "int64",
        arrow_schema::DataType::Float64 => "float64",
        arrow_schema::DataType::Utf8 => "utf8",
        arrow_schema::DataType::LargeUtf8 => "large_utf8",
        arrow_schema::DataType::Binary => "binary",
        arrow_schema::DataType::LargeBinary => "large_binary",
        arrow_schema::DataType::List(_) => "list",
        arrow_schema::DataType::LargeList(_) => "large_list",
        arrow_schema::DataType::Struct(_) => "struct",
        arrow_schema::DataType::Timestamp(_, _) => "timestamp",
        arrow_schema::DataType::Decimal128(_, _) => "decimal128",
        arrow_schema::DataType::Decimal256(_, _) => "decimal256",
        arrow_schema::DataType::Dictionary(_, _) => "dictionary",
        _ => "unknown",
    };
    rustler::types::atom::Atom::from_str(env, s).unwrap()
}

/// Return the schema ref of a record batch.
#[rustler::nif]
pub fn record_batch_schema(batch: ResourceArc<ExArrowRecordBatch>) -> rustler::resource::ResourceArc<ExArrowSchema> {
    let schema_handle = ExArrowSchema {
        schema: batch.batch.schema().clone(),
    };
    ResourceArc::new(schema_handle)
}

/// Return the number of rows in a record batch.
/// Returns as u64 so the value is always non-negative and never overflows i64.
#[rustler::nif]
pub fn record_batch_num_rows(batch: ResourceArc<ExArrowRecordBatch>) -> u64 {
    batch.batch.num_rows() as u64
}

/// Return the schema of an IPC stream (without consuming it).
#[rustler::nif]
pub fn ipc_stream_schema<'a>(env: Env<'a>, stream: ResourceArc<ExArrowIpcStream>) -> Term<'a> {
    let schema_ref = match &stream.reader {
        IpcStreamBacking::Binary(m) => {
            let guard = match m.lock() {
                Ok(g) => g,
                Err(_) => return err_encode(env, "stream lock"),
            };
            guard.schema()
        }
        IpcStreamBacking::File(m) => {
            let guard = match m.lock() {
                Ok(g) => g,
                Err(_) => return err_encode(env, "stream lock"),
            };
            guard.schema()
        }
    };
    let schema_handle = ExArrowSchema {
        schema: schema_ref,
    };
    ResourceArc::new(schema_handle).encode(env)
}

/// Read the next record batch from the stream. Returns {:ok, batch_ref} or :done or {:error, msg}.
#[rustler::nif]
pub fn ipc_stream_next<'a>(env: Env<'a>, stream: ResourceArc<ExArrowIpcStream>) -> Term<'a> {
    let next_result = match &stream.reader {
        IpcStreamBacking::Binary(m) => {
            let mut guard = match m.lock() {
                Ok(g) => g,
                Err(_) => return err_encode(env, "stream lock"),
            };
            guard.next()
        }
        IpcStreamBacking::File(m) => {
            let mut guard = match m.lock() {
                Ok(g) => g,
                Err(_) => return err_encode(env, "stream lock"),
            };
            guard.next()
        }
    };
    match next_result {
        None => rustler::types::atom::Atom::from_str(env, "done").unwrap().encode(env),
        Some(Err(e)) => err_encode(env, &e.to_string()),
        Some(Ok(batch)) => {
            let handle = ExArrowRecordBatch { batch };
            ok_encode(env, ResourceArc::new(handle))
        }
    }
}

/// Write schema and record batches to binary. Returns {:ok, binary} or {:error, msg}.
#[rustler::nif]
pub fn ipc_writer_to_binary<'a>(
    env: Env<'a>,
    schema: ResourceArc<ExArrowSchema>,
    batches: Vec<ResourceArc<ExArrowRecordBatch>>,
) -> Term<'a> {
    let mut buf = Vec::new();
    let mut writer = match StreamWriter::try_new(&mut buf, schema.schema.as_ref()) {
        Ok(w) => w,
        Err(e) => return err_encode(env, &e.to_string()),
    };
    for batch_ref in &batches {
        if let Err(e) = writer.write(&batch_ref.batch) {
            return err_encode(env, &e.to_string());
        }
    }
    if let Err(e) = writer.finish() {
        return err_encode(env, &e.to_string());
    }
    let mut owned = match rustler::OwnedBinary::new(buf.len()) {
        Some(b) => b,
        None => return err_encode(env, "binary alloc"),
    };
    owned.as_mut_slice().copy_from_slice(&buf);
    let binary = rustler::Binary::from_owned(owned, env);
    ok_encode(env, binary)
}

/// Open IPC file format (random access) from path. Returns {:ok, file_ref} or {:error, msg}.
#[rustler::nif]
pub fn ipc_file_open<'a>(env: Env<'a>, path: String) -> Term<'a> {
    let file = match std::fs::File::open(&path) {
        Ok(f) => f,
        Err(e) => return err_encode(env, &e.to_string()),
    };
    let reader = match FileReader::try_new_buffered(file, None) {
        Ok(r) => r,
        Err(e) => return err_encode(env, &e.to_string()),
    };
    let ipc_file = ExArrowIpcFile {
        backing: IpcFileBacking::File(std::sync::Mutex::new(reader)),
    };
    ok_encode(env, ResourceArc::new(ipc_file))
}

/// Open IPC file format from in-memory binary (random access). Returns {:ok, file_ref} or {:error, msg}.
#[rustler::nif]
pub fn ipc_file_open_from_binary<'a>(env: Env<'a>, binary: rustler::Binary) -> Term<'a> {
    let bytes = binary.as_slice().to_vec();
    let cursor = Cursor::new(bytes);
    let reader = match FileReader::try_new(cursor, None) {
        Ok(r) => r,
        Err(e) => return err_encode(env, &e.to_string()),
    };
    let ipc_file = ExArrowIpcFile {
        backing: IpcFileBacking::Binary(std::sync::Mutex::new(reader)),
    };
    ok_encode(env, ResourceArc::new(ipc_file))
}

/// Return the schema of an IPC file.
#[rustler::nif]
pub fn ipc_file_schema<'a>(env: Env<'a>, file: ResourceArc<ExArrowIpcFile>) -> Term<'a> {
    let schema_ref = match &file.backing {
        IpcFileBacking::File(m) => match m.lock() {
            Ok(g) => g.schema(),
            Err(_) => return err_encode(env, "file lock"),
        },
        IpcFileBacking::Binary(m) => match m.lock() {
            Ok(g) => g.schema(),
            Err(_) => return err_encode(env, "file lock"),
        },
    };
    let schema_handle = ExArrowSchema {
        schema: schema_ref,
    };
    ResourceArc::new(schema_handle).encode(env)
}

/// Return the number of record batches in an IPC file.
#[rustler::nif]
pub fn ipc_file_num_batches(file: ResourceArc<ExArrowIpcFile>) -> u64 {
    match &file.backing {
        IpcFileBacking::File(m) => match m.lock() {
            Ok(g) => g.num_batches() as u64,
            Err(_) => 0,
        },
        IpcFileBacking::Binary(m) => match m.lock() {
            Ok(g) => g.num_batches() as u64,
            Err(_) => 0,
        },
    }
}

/// Get the record batch at the given index (0-based). Returns {:ok, batch_ref} or {:error, msg}.
#[rustler::nif]
pub fn ipc_file_get_batch<'a>(
    env: Env<'a>,
    file: ResourceArc<ExArrowIpcFile>,
    index: u64,
) -> Term<'a> {
    let index = index as usize;
    let batch_result = match &file.backing {
        IpcFileBacking::File(m) => {
            let mut guard = match m.lock() {
                Ok(g) => g,
                Err(_) => return err_encode(env, "file lock"),
            };
            if let Err(e) = guard.set_index(index) {
                return err_encode(env, &e.to_string());
            }
            guard.next()
        }
        IpcFileBacking::Binary(m) => {
            let mut guard = match m.lock() {
                Ok(g) => g,
                Err(_) => return err_encode(env, "file lock"),
            };
            if let Err(e) = guard.set_index(index) {
                return err_encode(env, &e.to_string());
            }
            guard.next()
        }
    };
    match batch_result {
        None => err_encode(env, "batch index out of range"),
        Some(Err(e)) => err_encode(env, &e.to_string()),
        Some(Ok(batch)) => {
            let handle = ExArrowRecordBatch { batch };
            ok_encode(env, ResourceArc::new(handle))
        }
    }
}

/// Write schema and record batches to file (stream format). Returns :ok or {:error, msg}.
#[rustler::nif]
pub fn ipc_writer_to_file<'a>(
    env: Env<'a>,
    path: String,
    schema: ResourceArc<ExArrowSchema>,
    batches: Vec<ResourceArc<ExArrowRecordBatch>>,
) -> Term<'a> {
    let file = match std::fs::File::create(&path) {
        Ok(f) => f,
        Err(e) => return err_encode(env, &e.to_string()),
    };
    let mut writer = match StreamWriter::try_new(file, schema.schema.as_ref()) {
        Ok(w) => w,
        Err(e) => return err_encode(env, &e.to_string()),
    };
    for batch_ref in &batches {
        if let Err(e) = writer.write(&batch_ref.batch) {
            return err_encode(env, &e.to_string());
        }
    }
    if let Err(e) = writer.finish() {
        return err_encode(env, &e.to_string());
    }
    rustler::types::atom::Atom::from_str(env, "ok").unwrap().encode(env)
}

// ── Nx support NIFs ──────────────────────────────────────────────────────────

/// Extract a numeric column from a record batch as raw little-endian bytes for `Nx.from_binary/2`.
///
/// Returns `{:ok, {binary, type_string, num_rows}}` or `{:error, msg}`.
///
/// `type_string` is one of `"s8"`, `"s16"`, `"s32"`, `"s64"`, `"u8"`, `"u16"`,
/// `"u32"`, `"u64"`, `"f32"`, `"f64"`.  Float columns use IEEE 754 native byte order.
///
/// Null/validity bitmaps are ignored — null positions are returned as zero bytes.
#[rustler::nif]
pub fn record_batch_column_buffer<'a>(
    env: Env<'a>,
    batch: ResourceArc<ExArrowRecordBatch>,
    col_name: String,
) -> Term<'a> {
    let schema = batch.batch.schema();
    let col_idx = match schema.index_of(&col_name) {
        Ok(i) => i,
        Err(_) => return err_encode(env, &format!("column '{}' not found", col_name)),
    };
    let array = batch.batch.column(col_idx);
    extract_primitive_buffer(env, array)
}

macro_rules! primitive_buffer {
    ($env:expr, $array:expr, $ArrowType:ty, $dtype_str:expr) => {{
        let arr: &PrimitiveArray<$ArrowType> = match $array.as_any().downcast_ref() {
            Some(a) => a,
            None => return err_encode($env, "internal: type downcast failed"),
        };
        let offset = arr.offset();
        let len = arr.len();
        let values = arr.values();
        let slice = &values[offset..offset + len];
        let byte_size =
            std::mem::size_of::<<$ArrowType as arrow_array::types::ArrowPrimitiveType>::Native>();
        // SAFETY: slice is a valid aligned slice of a numeric primitive type whose
        // in-memory representation is exactly `len * byte_size` bytes.
        let bytes: &[u8] = unsafe {
            std::slice::from_raw_parts(slice.as_ptr() as *const u8, len * byte_size)
        };
        let mut owned = match rustler::OwnedBinary::new(bytes.len()) {
            Some(b) => b,
            None => return err_encode($env, "binary alloc"),
        };
        owned.as_mut_slice().copy_from_slice(bytes);
        let binary = rustler::Binary::from_owned(owned, $env);
        ok_encode($env, (binary, $dtype_str, len as u64))
    }};
}

fn extract_primitive_buffer<'a>(env: Env<'a>, array: &ArrayRef) -> Term<'a> {
    match array.data_type() {
        DataType::Int8 => primitive_buffer!(env, array, Int8Type, "s8"),
        DataType::Int16 => primitive_buffer!(env, array, Int16Type, "s16"),
        DataType::Int32 => primitive_buffer!(env, array, Int32Type, "s32"),
        DataType::Int64 => primitive_buffer!(env, array, Int64Type, "s64"),
        DataType::UInt8 => primitive_buffer!(env, array, UInt8Type, "u8"),
        DataType::UInt16 => primitive_buffer!(env, array, UInt16Type, "u16"),
        DataType::UInt32 => primitive_buffer!(env, array, UInt32Type, "u32"),
        DataType::UInt64 => primitive_buffer!(env, array, UInt64Type, "u64"),
        DataType::Float32 => primitive_buffer!(env, array, Float32Type, "f32"),
        DataType::Float64 => primitive_buffer!(env, array, Float64Type, "f64"),
        dt => err_encode(env, &format!("unsupported column type for Nx: {:?}", dt)),
    }
}

/// Create a single-column `RecordBatch` from raw bytes (the reverse of `record_batch_column_buffer`).
///
/// `dtype_str` must be one of `"s8"`, `"s16"`, `"s32"`, `"s64"`, `"u8"`, `"u16"`,
/// `"u32"`, `"u64"`, `"f32"`, `"f64"`.
///
/// Returns `{:ok, batch_ref}` or `{:error, msg}`.
#[rustler::nif]
pub fn record_batch_from_column_binary<'a>(
    env: Env<'a>,
    col_name: String,
    binary: rustler::Binary,
    dtype_str: String,
    length: u64,
) -> Term<'a> {
    let length = length as usize;
    let bytes = binary.as_slice();

    macro_rules! build_primitive_batch {
        ($NativeType:ty, $ArrowDType:expr) => {{
            let elem_size = std::mem::size_of::<$NativeType>();
            if bytes.len() != length * elem_size {
                return err_encode(
                    env,
                    &format!(
                        "binary length mismatch: expected {} bytes ({} * {}), got {}",
                        length * elem_size,
                        length,
                        elem_size,
                        bytes.len()
                    ),
                );
            }
            let buf = Buffer::from_slice_ref(bytes);
            let array_data = match ArrayData::builder($ArrowDType)
                .len(length)
                .add_buffer(buf)
                .build()
            {
                Ok(d) => d,
                Err(e) => return err_encode(env, &e.to_string()),
            };
            let array: ArrayRef = Arc::new(PrimitiveArray::<
                <$NativeType as NativeTypeAlias>::ArrowType,
            >::from(array_data));
            array
        }};
    }

    let array: ArrayRef = match dtype_str.as_str() {
        "s8" => build_primitive_batch!(i8, DataType::Int8),
        "s16" => build_primitive_batch!(i16, DataType::Int16),
        "s32" => build_primitive_batch!(i32, DataType::Int32),
        "s64" => build_primitive_batch!(i64, DataType::Int64),
        "u8" => build_primitive_batch!(u8, DataType::UInt8),
        "u16" => build_primitive_batch!(u16, DataType::UInt16),
        "u32" => build_primitive_batch!(u32, DataType::UInt32),
        "u64" => build_primitive_batch!(u64, DataType::UInt64),
        "f32" => build_primitive_batch!(f32, DataType::Float32),
        "f64" => build_primitive_batch!(f64, DataType::Float64),
        other => return err_encode(env, &format!("unknown dtype '{}' for column creation", other)),
    };

    let schema = Arc::new(Schema::new(vec![Field::new(
        &col_name,
        array.data_type().clone(),
        false,
    )]));
    match RecordBatch::try_new(schema, vec![array]) {
        Ok(batch) => ok_encode(env, ResourceArc::new(ExArrowRecordBatch { batch })),
        Err(e) => err_encode(env, &e.to_string()),
    }
}

// Helper trait to associate Rust native types with Arrow type markers.
trait NativeTypeAlias {
    type ArrowType: arrow_array::types::ArrowPrimitiveType;
}
impl NativeTypeAlias for i8 {
    type ArrowType = Int8Type;
}
impl NativeTypeAlias for i16 {
    type ArrowType = Int16Type;
}
impl NativeTypeAlias for i32 {
    type ArrowType = Int32Type;
}
impl NativeTypeAlias for i64 {
    type ArrowType = Int64Type;
}
impl NativeTypeAlias for u8 {
    type ArrowType = UInt8Type;
}
impl NativeTypeAlias for u16 {
    type ArrowType = UInt16Type;
}
impl NativeTypeAlias for u32 {
    type ArrowType = UInt32Type;
}
impl NativeTypeAlias for u64 {
    type ArrowType = UInt64Type;
}
impl NativeTypeAlias for f32 {
    type ArrowType = Float32Type;
}
impl NativeTypeAlias for f64 {
    type ArrowType = Float64Type;
}

/// Write schema and record batches in IPC file format (random-access footer). Returns :ok or {:error, msg}.
#[rustler::nif]
pub fn ipc_file_writer_to_file<'a>(
    env: Env<'a>,
    path: String,
    schema: ResourceArc<ExArrowSchema>,
    batches: Vec<ResourceArc<ExArrowRecordBatch>>,
) -> Term<'a> {
    let file = match std::fs::File::create(&path) {
        Ok(f) => f,
        Err(e) => return err_encode(env, &e.to_string()),
    };
    let mut writer = match FileWriter::try_new_buffered(file, schema.schema.as_ref()) {
        Ok(w) => w,
        Err(e) => return err_encode(env, &e.to_string()),
    };
    for batch_ref in &batches {
        if let Err(e) = writer.write(&batch_ref.batch) {
            return err_encode(env, &e.to_string());
        }
    }
    if let Err(e) = writer.finish() {
        return err_encode(env, &e.to_string());
    }
    rustler::types::atom::Atom::from_str(env, "ok").unwrap().encode(env)
}