graphar-flight 0.1.2

use arrow_array::RecordBatch;
use arrow_schema::SchemaRef;
use redis::aio::MultiplexedConnection;

use crate::{
    convert::{response_to_batch, response_to_batch_auto},
    error::Result,
};

/// Executes Cypher queries against FalkorDB and converts results to Arrow.
///
/// `MultiplexedConnection` is `Clone + Send + Sync`, so this struct is cheap
/// to clone per-request — no mutex needed.
#[derive(Clone)]
pub struct FalkorExecutor {
    conn: MultiplexedConnection,
    graph: String,
}

impl FalkorExecutor {
    pub async fn connect(url: &str, graph: impl Into<String>) -> Result<Self> {
        let client = redis::Client::open(url)?;
        let conn = client.get_multiplexed_async_connection().await?;
        Ok(Self {
            conn,
            graph: graph.into(),
        })
    }

    /// Run a read Cypher query; convert the response using the given schema.
    pub async fn query(&mut self, cypher: &str, schema: &SchemaRef) -> Result<RecordBatch> {
        let raw: redis::Value = redis::cmd("GRAPH.QUERY")
            .arg(&self.graph)
            .arg(cypher)
            .query_async(&mut self.conn)
            .await?;
        response_to_batch(raw, schema)
    }

    /// Run a read Cypher query using the auto-string schema fallback.
    ///
    /// Column names are taken from the FalkorDB response header; all values
    /// are returned as `Utf8`. No schema pre-registration is required.
    pub async fn query_auto(&mut self, cypher: &str) -> Result<RecordBatch> {
        let raw: redis::Value = redis::cmd("GRAPH.QUERY")
            .arg(&self.graph)
            .arg(cypher)
            .query_async(&mut self.conn)
            .await?;
        response_to_batch_auto(raw)
    }

    /// Run a write Cypher query (CREATE / MERGE / DELETE …) with no result rows.
    /// Returns the number of nodes/relationships affected (from the stats line).
    pub async fn execute(&mut self, cypher: &str) -> Result<u64> {
        let raw: redis::Value = redis::cmd("GRAPH.QUERY")
            .arg(&self.graph)
            .arg(cypher)
            .query_async(&mut self.conn)
            .await?;
        Ok(parse_affected_count(raw))
    }

    pub fn graph(&self) -> &str {
        &self.graph
    }
}

/// Pull a rough "nodes created + relationships created" count from the stats
/// array that FalkorDB appends to every response.
fn parse_affected_count(response: redis::Value) -> u64 {
    let outer = match response {
        redis::Value::Array(v) => v,
        _ => return 0,
    };
    // Stats are in the last element.
    let stats = match outer.last() {
        Some(redis::Value::Array(s)) => s,
        _ => return 0,
    };
    let mut total = 0u64;
    for stat in stats {
        if let redis::Value::BulkString(b) = stat {
            let s = String::from_utf8_lossy(b);
            // e.g. "Nodes created: 3" or "Relationships created: 1"
            if (s.contains("created:") || s.contains("deleted:"))
                && let Some(n) = s
                    .split(':')
                    .nth(1)
                    .and_then(|p| p.trim().parse::<u64>().ok())
            {
                total += n;
            }
        }
    }
    total
}