//! Conversions between PropertyGraph and Arrow RecordBatch.
//!
//! Provides zero-copy-friendly conversions used by both the Flight server
//! (serialising graph data to send) and the Flight client (deserialising
//! received batches back into graph data).

#[cfg(feature = "distributed")]
use std::sync::Arc;

#[cfg(feature = "distributed")]
use arrow::array::{
    ArrayRef, ListBuilder, RecordBatch, StringBuilder, StringArray, UInt64Array,
};
#[cfg(feature = "distributed")]
use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
#[cfg(feature = "distributed")]
use arrow::error::ArrowError;

#[cfg(feature = "distributed")]
use crate::error::{CypherError, CypherResult, ExecutionError};
#[cfg(feature = "distributed")]
use crate::graph::{EdgeLike, NodeLike};

// ── Schemas ──────────────────────────────────────────────────────────────────

/// Arrow schema for a nodes record batch.
#[cfg(feature = "distributed")]
pub fn nodes_schema() -> SchemaRef {
    Arc::new(Schema::new(vec![
        Field::new("vertex_id", DataType::UInt64, false),
        Field::new(
            "labels",
            DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))),
            true,
        ),
        Field::new("properties_json", DataType::Utf8, true),
    ]))
}

/// Arrow schema for an edges record batch.
#[cfg(feature = "distributed")]
pub fn edges_schema() -> SchemaRef {
    Arc::new(Schema::new(vec![
        Field::new("edge_id", DataType::UInt64, false),
        Field::new("src_id", DataType::UInt64, false),
        Field::new("dst_id", DataType::UInt64, false),
        Field::new("rel_type", DataType::Utf8, false),
        Field::new("properties_json", DataType::Utf8, true),
    ]))
}

// ── Encode ───────────────────────────────────────────────────────────────────

/// Convert a slice of node references to a single `RecordBatch`.
///
/// Accepts any type implementing `NodeLike` so that different backends
/// (e.g. `NKNodeRef`, `&GraphNode`) can both be serialised without copying.
#[cfg(feature = "distributed")]
pub fn nodes_to_record_batch<N: NodeLike>(nodes: &[N]) -> CypherResult<RecordBatch> {
    let schema = nodes_schema();

    if nodes.is_empty() {
        return RecordBatch::try_new(
            schema,
            vec![
                Arc::new(UInt64Array::from(Vec::<u64>::new())) as ArrayRef,
                Arc::new(ListBuilder::new(StringBuilder::new()).finish()) as ArrayRef,
                Arc::new(StringArray::from(Vec::<String>::new())) as ArrayRef,
            ],
        )
        .map_err(arrow_err);
    }

    // vertex_id column
    let vertex_ids: Vec<u64> = nodes.iter().map(|n| n.id()).collect();

    // labels column  (List<Utf8>)
    let mut labels_builder =
        ListBuilder::new(StringBuilder::new())
            .with_field(Arc::new(Field::new("item", DataType::Utf8, true)));
    for node in nodes {
        for label in node.labels() {
            labels_builder.values().append_value(label);
        }
        labels_builder.append(true);
    }

    // properties_json column
    let props_json: Vec<String> = nodes
        .iter()
        .map(|n| serde_json::to_string(n.properties()).unwrap_or_default())
        .collect();

    RecordBatch::try_new(
        schema,
        vec![
            Arc::new(UInt64Array::from(vertex_ids)) as ArrayRef,
            Arc::new(labels_builder.finish()) as ArrayRef,
            Arc::new(StringArray::from(props_json)) as ArrayRef,
        ],
    )
    .map_err(arrow_err)
}

/// Convert a slice of edge references (with their endpoints) to a single `RecordBatch`.
///
/// Accepts any type implementing `EdgeLike` so that different backends
/// (e.g. `NKEdgeRef`, `&GraphEdge`) can both be serialised without copying.
#[cfg(feature = "distributed")]
pub fn edges_to_record_batch<E: EdgeLike>(
    edges: &[E],
    endpoints: &[(u64, u64)],
) -> CypherResult<RecordBatch> {
    let schema = edges_schema();

    if edges.is_empty() {
        return RecordBatch::try_new(
            schema,
            vec![
                Arc::new(UInt64Array::from(Vec::<u64>::new())) as ArrayRef,
                Arc::new(UInt64Array::from(Vec::<u64>::new())) as ArrayRef,
                Arc::new(UInt64Array::from(Vec::<u64>::new())) as ArrayRef,
                Arc::new(StringArray::from(Vec::<String>::new())) as ArrayRef,
                Arc::new(StringArray::from(Vec::<String>::new())) as ArrayRef,
            ],
        )
        .map_err(arrow_err);
    }

    let edge_ids: Vec<u64> = edges.iter().map(|e| e.id()).collect();
    let src_ids: Vec<u64> = endpoints.iter().map(|(s, _)| *s).collect();
    let dst_ids: Vec<u64> = endpoints.iter().map(|(_, d)| *d).collect();
    let rel_types: Vec<&str> = edges.iter().map(|e| e.rel_type()).collect();
    let props_json: Vec<String> = edges
        .iter()
        .map(|e| serde_json::to_string(e.properties()).unwrap_or_default())
        .collect();

    RecordBatch::try_new(
        schema,
        vec![
            Arc::new(UInt64Array::from(edge_ids)) as ArrayRef,
            Arc::new(UInt64Array::from(src_ids)) as ArrayRef,
            Arc::new(UInt64Array::from(dst_ids)) as ArrayRef,
            Arc::new(StringArray::from(rel_types)) as ArrayRef,
            Arc::new(StringArray::from(props_json)) as ArrayRef,
        ],
    )
    .map_err(arrow_err)
}

// ── Helpers ───────────────────────────────────────────────────────────────────

#[cfg(feature = "distributed")]
fn arrow_err(e: ArrowError) -> CypherError {
    CypherError::Execution(ExecutionError::Internal(e.to_string()))
}

// ── Tests ─────────────────────────────────────────────────────────────────────

#[cfg(all(test, feature = "distributed"))]
mod tests {
    use super::*;
    use crate::graph::{GraphEdge, GraphNode, PropertyGraph};
    use indexmap::IndexMap;

    #[test]
    fn test_empty_nodes() {
        let batch = nodes_to_record_batch::<&GraphNode>(&[]).unwrap();
        assert_eq!(batch.num_rows(), 0);
        assert_eq!(batch.num_columns(), 3);
    }

    #[test]
    fn test_nodes_round_trip() {
        let mut graph = PropertyGraph::new();
        let mut props = IndexMap::new();
        props.insert(
            "name".to_string(),
            crate::graph::PropertyValue::String("Alice".to_string()),
        );
        graph.create_node(vec!["Person"], props);
        graph.create_node(vec!["Person", "Engineer"], IndexMap::new());

        let nodes: Vec<&GraphNode> = graph.all_nodes().collect();
        let batch = nodes_to_record_batch(&nodes).unwrap();

        assert_eq!(batch.num_rows(), 2);
        assert_eq!(batch.schema().field(0).name(), "vertex_id");
        assert_eq!(batch.schema().field(1).name(), "labels");
        assert_eq!(batch.schema().field(2).name(), "properties_json");
    }

    #[test]
    fn test_empty_edges() {
        let batch = edges_to_record_batch::<&GraphEdge>(&[], &[]).unwrap();
        assert_eq!(batch.num_rows(), 0);
        assert_eq!(batch.num_columns(), 5);
    }

    #[test]
    fn test_edges_round_trip() {
        let mut graph = PropertyGraph::new();
        let n1 = graph.create_node(vec!["Person"], IndexMap::new()).id;
        let n2 = graph.create_node(vec!["Person"], IndexMap::new()).id;
        graph
            .create_relationship(n1, n2, "KNOWS", IndexMap::new())
            .unwrap();

        let edges: Vec<&GraphEdge> = graph.all_edges().collect();
        let endpoints: Vec<(u64, u64)> = edges
            .iter()
            .filter_map(|e| graph.get_edge_endpoints(e.id()))
            .collect();

        let batch = edges_to_record_batch(&edges, &endpoints).unwrap();
        assert_eq!(batch.num_rows(), 1);
        assert_eq!(batch.schema().field(3).name(), "rel_type");
    }
}
