//! Arrow Flight RPC client (arrow-flight 57.x / tonic 0.14).
//!
//! Used by other pods / the query router to fetch graph data from a remote
//! partition over zero-copy Arrow Flight.

#[cfg(feature = "distributed")]
use arrow::record_batch::RecordBatch;
#[cfg(feature = "distributed")]
use arrow_flight::{
    decode::FlightRecordBatchStream,
    flight_service_client::FlightServiceClient,
    Action, Ticket,
};
#[cfg(feature = "distributed")]
use futures::{StreamExt, TryStreamExt};
#[cfg(feature = "distributed")]
use tonic::transport::Channel;

#[cfg(feature = "distributed")]
use crate::distributed::network::superstep::{SuperstepRequest, SuperstepResponse};
#[cfg(feature = "distributed")]
use crate::error::{CypherError, CypherResult, ExecutionError};

// ── Client ────────────────────────────────────────────────────────────────────

/// Client for a remote graph partition's Flight service.
#[cfg(feature = "distributed")]
pub struct GraphFlightClient {
    inner: FlightServiceClient<Channel>,
}

#[cfg(feature = "distributed")]
impl GraphFlightClient {
    /// Connect to a remote partition's Flight server.
    ///
    /// `endpoint` examples:
    ///   - `"http://ocg-partition-2:8815"`  — plain HTTP
    ///   - `"https://ocg-partition-2:8815"` — TLS; uses system CA roots
    ///
    /// For TLS with a custom CA, set `OCG_TLS_CA` to the path of the PEM CA
    /// certificate file before calling `connect`.
    pub async fn connect(endpoint: impl Into<String>) -> CypherResult<Self> {
        let ep = endpoint.into();
        let use_tls = ep.starts_with("https://");

        let channel = if use_tls {
            let mut tls_config = tonic::transport::ClientTlsConfig::new();

            // Optional custom CA cert (e.g. self-signed cluster CA)
            if let Ok(ca_path) = std::env::var("OCG_TLS_CA") {
                let ca_pem = std::fs::read(&ca_path)
                    .map_err(|e| internal(format!("read TLS CA {ca_path}: {e}")))?;
                let ca = tonic::transport::Certificate::from_pem(ca_pem);
                tls_config = tls_config.ca_certificate(ca);
            }

            // Optional client cert for mutual TLS
            if let (Ok(cert_path), Ok(key_path)) = (
                std::env::var("OCG_TLS_CERT"),
                std::env::var("OCG_TLS_KEY"),
            ) {
                let cert = std::fs::read(&cert_path)
                    .map_err(|e| internal(format!("read client cert {cert_path}: {e}")))?;
                let key = std::fs::read(&key_path)
                    .map_err(|e| internal(format!("read client key {key_path}: {e}")))?;
                let identity = tonic::transport::Identity::from_pem(cert, key);
                tls_config = tls_config.identity(identity);
            }

            Channel::from_shared(ep.clone())
                .map_err(|e| internal(format!("invalid endpoint {ep}: {e}")))?
                .tls_config(tls_config)
                .map_err(|e| internal(format!("TLS config for {ep}: {e}")))?
                .connect()
                .await
                .map_err(|e| internal(format!("connect (TLS) to {ep}: {e}")))?
        } else {
            Channel::from_shared(ep.clone())
                .map_err(|e| internal(format!("invalid endpoint {ep}: {e}")))?
                .connect()
                .await
                .map_err(|e| internal(format!("connect to {ep}: {e}")))?
        };

        let inner = FlightServiceClient::new(channel);
        Ok(Self { inner })
    }

    // ── Actions ───────────────────────────────────────────────────────────

    /// Ping the remote server. Returns `"OK"` on success.
    pub async fn health_check(&mut self) -> CypherResult<String> {
        let action = Action { r#type: "HealthCheck".into(), body: vec![].into() };
        let mut stream = self
            .inner
            .do_action(action)
            .await
            .map_err(|e| internal(format!("HealthCheck RPC: {e}")))?
            .into_inner();

        let msg = stream
            .message()
            .await
            .map_err(|e| internal(format!("HealthCheck recv: {e}")))?
            .ok_or_else(|| internal("HealthCheck: empty response"))?;

        String::from_utf8(msg.body.to_vec())
            .map_err(|e| internal(format!("HealthCheck UTF-8: {e}")))
    }

    /// Return the partition ID managed by the remote server.
    pub async fn get_partition_id(&mut self) -> CypherResult<u32> {
        let action = Action { r#type: "GetPartitionId".into(), body: vec![].into() };
        let mut stream = self
            .inner
            .do_action(action)
            .await
            .map_err(|e| internal(format!("GetPartitionId RPC: {e}")))?
            .into_inner();

        let msg = stream
            .message()
            .await
            .map_err(|e| internal(format!("GetPartitionId recv: {e}")))?
            .ok_or_else(|| internal("GetPartitionId: empty response"))?;

        let s = String::from_utf8(msg.body.to_vec())
            .map_err(|e| internal(format!("GetPartitionId UTF-8: {e}")))?;
        s.parse::<u32>()
            .map_err(|e| internal(format!("GetPartitionId parse: {e}")))
    }

    /// Instruct the remote server to load a partition from its snapshot.
    ///
    /// Sends a `LoadPartition` DoAction with JSON body containing the
    /// `storage_path`. The server restores from Parquet + WAL replay.
    ///
    /// Returns the server's response message on success.
    pub async fn load_partition(
        &mut self,
        storage_path: &str,
    ) -> CypherResult<String> {
        let body = serde_json::json!({ "storage_path": storage_path });
        let action = Action {
            r#type: "LoadPartition".into(),
            body: serde_json::to_vec(&body)
                .map_err(|e| internal(format!("LoadPartition serialize: {e}")))?
                .into(),
        };
        let mut stream = self
            .inner
            .do_action(action)
            .await
            .map_err(|e| internal(format!("LoadPartition RPC: {e}")))?
            .into_inner();

        let msg = stream
            .message()
            .await
            .map_err(|e| internal(format!("LoadPartition recv: {e}")))?
            .ok_or_else(|| internal("LoadPartition: empty response"))?;

        String::from_utf8(msg.body.to_vec())
            .map_err(|e| internal(format!("LoadPartition UTF-8: {e}")))
    }

    // ── BSP superstep ─────────────────────────────────────────────────────

    /// Send a BSP superstep request to the remote partition and return its response.
    ///
    /// The request is serialised to JSON and sent as a `RunSuperstep` DoAction.
    /// The response is deserialised from the action result body.
    pub async fn run_superstep(
        &mut self,
        request: &SuperstepRequest,
    ) -> CypherResult<SuperstepResponse> {
        let body = serde_json::to_vec(request)
            .map_err(|e| internal(format!("RunSuperstep serialize: {e}")))?;
        let action = Action { r#type: "RunSuperstep".into(), body: body.into() };

        let mut stream = self
            .inner
            .do_action(action)
            .await
            .map_err(|e| internal(format!("RunSuperstep RPC: {e}")))?
            .into_inner();

        let msg = stream
            .message()
            .await
            .map_err(|e| internal(format!("RunSuperstep recv: {e}")))?
            .ok_or_else(|| internal("RunSuperstep: empty response"))?;

        serde_json::from_slice(&msg.body)
            .map_err(|e| internal(format!("RunSuperstep deserialize: {e}")))
    }

    // ── Data fetches ──────────────────────────────────────────────────────

    /// Fetch all nodes from the remote partition.
    ///
    /// Returns one `RecordBatch` per IPC message in the response stream —
    /// typically one batch unless the partition is very large.
    pub async fn fetch_nodes(&mut self) -> CypherResult<Vec<RecordBatch>> {
        self.fetch_ticket(b"nodes").await
    }

    /// Fetch all edges from the remote partition.
    pub async fn fetch_edges(&mut self) -> CypherResult<Vec<RecordBatch>> {
        self.fetch_ticket(b"edges").await
    }

    // ── Internal ──────────────────────────────────────────────────────────

    async fn fetch_ticket(&mut self, ticket: &[u8]) -> CypherResult<Vec<RecordBatch>> {
        let tkt = Ticket { ticket: ticket.to_vec().into() };

        let response = self
            .inner
            .do_get(tkt)
            .await
            .map_err(|e| internal(format!("do_get RPC: {e}")))?
            .into_inner();

        // FlightRecordBatchStream::new_from_flight_data decodes the
        // Arrow IPC schema + data messages without any extra copies.
        let mut batch_stream = FlightRecordBatchStream::new_from_flight_data(
            response.map_err(|e| arrow_flight::error::FlightError::Tonic(Box::new(e))),
        );

        let mut batches = Vec::new();
        while let Some(result) = batch_stream.next().await {
            let batch = result
                .map_err(|e| internal(format!("decode batch: {e}")))?;
            batches.push(batch);
        }
        Ok(batches)
    }
}

// ── Error helpers ─────────────────────────────────────────────────────────────

#[cfg(feature = "distributed")]
fn internal(msg: impl Into<String>) -> CypherError {
    CypherError::Execution(ExecutionError::Internal(msg.into()))
}

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

#[cfg(all(test, feature = "distributed"))]
mod tests {
    // Integration tests require a live server; mark them ignored so they
    // don't run in CI without the infrastructure.

    #[tokio::test]
    #[ignore = "requires a running Flight server on localhost:8815"]
    async fn test_health_check_live() {
        use super::GraphFlightClient;
        let mut client = GraphFlightClient::connect("http://localhost:8815")
            .await
            .unwrap();
        assert_eq!(client.health_check().await.unwrap(), "OK");
    }

    #[tokio::test]
    #[ignore = "requires a running Flight server on localhost:8815"]
    async fn test_fetch_nodes_live() {
        use super::GraphFlightClient;
        let mut client = GraphFlightClient::connect("http://localhost:8815")
            .await
            .unwrap();
        let batches = client.fetch_nodes().await.unwrap();
        // Just assert we got at least one batch back
        assert!(!batches.is_empty());
    }
}
