//! Arrow Flight RPC server (arrow-flight 57.x / tonic 0.14).
//!
//! Runs inside each partition pod. Exposes nodes and edges of the local
//! partition over the Arrow Flight protocol for zero-copy transfer.

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

#[cfg(feature = "distributed")]
use arrow_flight::{
    encode::FlightDataEncoderBuilder,
    flight_service_server::{FlightService, FlightServiceServer},
    Action, ActionType, Criteria, Empty, FlightData, FlightDescriptor, FlightInfo,
    HandshakeRequest, HandshakeResponse, PollInfo, PutResult, SchemaResult, Ticket,
};
#[cfg(feature = "distributed")]
use futures::stream::BoxStream;
#[cfg(feature = "distributed")]
use futures::StreamExt;
#[cfg(feature = "distributed")]
use tokio::sync::RwLock;
#[cfg(feature = "distributed")]
use tonic::{Request, Response, Status, Streaming};

#[cfg(feature = "distributed")]
use crate::distributed::graph::PartitionedGraphBackend;
#[cfg(feature = "distributed")]
use crate::distributed::network::conversions::{
    edges_schema, edges_to_record_batch, nodes_schema, nodes_to_record_batch,
};
#[cfg(feature = "distributed")]
use crate::distributed::network::superstep::SuperstepRequest;
#[cfg(feature = "distributed")]
use crate::graph::NodeLike;
#[cfg(all(feature = "distributed", feature = "self-heal"))]
use crate::distributed::storage::snapshot_restore;
#[cfg(all(feature = "distributed", feature = "self-heal"))]
use crate::distributed::storage::wal::WalConfig;
#[cfg(feature = "distributed")]
use crate::graph::EdgeLike;
#[cfg(feature = "distributed")]
use crate::graph::GraphBackend;

// ── Service struct ────────────────────────────────────────────────────────────

/// Arrow Flight service that exposes a single graph partition.
#[cfg(feature = "distributed")]
pub struct GraphFlightService {
    backend: Arc<RwLock<PartitionedGraphBackend>>,
}

#[cfg(feature = "distributed")]
impl GraphFlightService {
    pub fn new(backend: Arc<RwLock<PartitionedGraphBackend>>) -> Self {
        Self { backend }
    }

    pub fn into_server(self) -> FlightServiceServer<Self> {
        FlightServiceServer::new(self)
    }
}

// ── FlightService impl ────────────────────────────────────────────────────────

#[cfg(feature = "distributed")]
#[tonic::async_trait]
impl FlightService for GraphFlightService {
    // All seven associated stream types use BoxStream<'static, _> — the
    // canonical pattern from the official arrow-flight examples.
    type HandshakeStream   = BoxStream<'static, Result<HandshakeResponse, Status>>;
    type ListFlightsStream = BoxStream<'static, Result<FlightInfo, Status>>;
    type DoGetStream       = BoxStream<'static, Result<FlightData, Status>>;
    type DoPutStream       = BoxStream<'static, Result<PutResult, Status>>;
    type DoExchangeStream  = BoxStream<'static, Result<FlightData, Status>>;
    type DoActionStream    = BoxStream<'static, Result<arrow_flight::Result, Status>>;
    type ListActionsStream = BoxStream<'static, Result<ActionType, Status>>;

    // ── Ticket handler ─────────────────────────────────────────────────────

    async fn do_get(
        &self,
        request: Request<Ticket>,
    ) -> Result<Response<Self::DoGetStream>, Status> {
        let ticket_bytes = request.into_inner().ticket;
        let ticket = String::from_utf8(ticket_bytes.to_vec())
            .map_err(|e| Status::invalid_argument(format!("Invalid ticket: {}", e)))?;

        match ticket.as_str() {
            "nodes" => self.stream_nodes().await,
            "edges" => self.stream_edges().await,
            other   => Err(Status::invalid_argument(format!("Unknown ticket: {}", other))),
        }
    }

    // ── Action handler ─────────────────────────────────────────────────────

    async fn do_action(
        &self,
        request: Request<Action>,
    ) -> Result<Response<Self::DoActionStream>, Status> {
        let action = request.into_inner();
        match action.r#type.as_str() {
            "HealthCheck"    => self.action_health_check().await,
            "GetPartitionId" => self.action_get_partition_id().await,
            "LoadPartition"  => self.action_load_partition(action.body.to_vec()).await,
            "RunSuperstep"   => self.action_run_superstep(action.body.to_vec()).await,
            other => Err(Status::invalid_argument(format!("Unknown action: {}", other))),
        }
    }

    // ── list_actions ───────────────────────────────────────────────────────

    async fn list_actions(
        &self,
        _request: Request<Empty>,
    ) -> Result<Response<Self::ListActionsStream>, Status> {
        let actions = vec![
            Ok(ActionType {
                r#type: "HealthCheck".to_string(),
                description: "Returns OK if the server is healthy".to_string(),
            }),
            Ok(ActionType {
                r#type: "GetPartitionId".to_string(),
                description: "Returns the partition ID this server manages".to_string(),
            }),
            Ok(ActionType {
                r#type: "LoadPartition".to_string(),
                description: "Load a partition from Parquet snapshot + WAL replay".to_string(),
            }),
            Ok(ActionType {
                r#type: "RunSuperstep".to_string(),
                description: "Execute one BSP superstep for a distributed graph algorithm".to_string(),
            }),
        ];
        Ok(Response::new(Box::pin(futures::stream::iter(actions))))
    }

    // ── Unimplemented stubs ────────────────────────────────────────────────

    async fn handshake(
        &self,
        _request: Request<Streaming<HandshakeRequest>>,
    ) -> Result<Response<Self::HandshakeStream>, Status> {
        Err(Status::unimplemented("handshake"))
    }

    async fn list_flights(
        &self,
        _request: Request<Criteria>,
    ) -> Result<Response<Self::ListFlightsStream>, Status> {
        Err(Status::unimplemented("list_flights"))
    }

    async fn get_flight_info(
        &self,
        _request: Request<FlightDescriptor>,
    ) -> Result<Response<FlightInfo>, Status> {
        Err(Status::unimplemented("get_flight_info"))
    }

    async fn poll_flight_info(
        &self,
        _request: Request<FlightDescriptor>,
    ) -> Result<Response<PollInfo>, Status> {
        Err(Status::unimplemented("poll_flight_info"))
    }

    async fn get_schema(
        &self,
        _request: Request<FlightDescriptor>,
    ) -> Result<Response<SchemaResult>, Status> {
        Err(Status::unimplemented("get_schema"))
    }

    async fn do_put(
        &self,
        _request: Request<Streaming<FlightData>>,
    ) -> Result<Response<Self::DoPutStream>, Status> {
        Err(Status::unimplemented("do_put"))
    }

    async fn do_exchange(
        &self,
        _request: Request<Streaming<FlightData>>,
    ) -> Result<Response<Self::DoExchangeStream>, Status> {
        Err(Status::unimplemented("do_exchange"))
    }
}

// ── Private helpers ───────────────────────────────────────────────────────────

#[cfg(feature = "distributed")]
impl GraphFlightService {
    async fn stream_nodes(&self) -> Result<Response<BoxStream<'static, Result<FlightData, Status>>>, Status> {
        let backend = self.backend.read().await;
        let nodes: Vec<_> = backend.all_nodes().collect();

        let batch = nodes_to_record_batch(&nodes)
            .map_err(|e| Status::internal(e.to_string()))?;

        let schema = nodes_schema();
        // Build as Vec<RecordBatch> → Stream<Item = Result<RecordBatch, ArrowError>>
        let batch_stream = futures::stream::once(
            async move { Ok::<_, arrow_flight::error::FlightError>(batch) }
        );

        let flight_stream = FlightDataEncoderBuilder::new()
            .with_schema(schema)
            .build(batch_stream)
            .map(|r| r.map_err(|e| Status::internal(e.to_string())));

        Ok(Response::new(Box::pin(flight_stream)))
    }

    async fn stream_edges(&self) -> Result<Response<BoxStream<'static, Result<FlightData, Status>>>, Status> {
        let backend = self.backend.read().await;
        let edges: Vec<_> = backend.all_edges().collect();
        let endpoints: Vec<(u64, u64)> = edges
            .iter()
            .filter_map(|e| backend.get_edge_endpoints(e.id()))
            .collect();

        let batch = edges_to_record_batch(&edges, &endpoints)
            .map_err(|e| Status::internal(e.to_string()))?;

        let schema = edges_schema();
        let batch_stream = futures::stream::once(
            async move { Ok::<_, arrow_flight::error::FlightError>(batch) }
        );

        let flight_stream = FlightDataEncoderBuilder::new()
            .with_schema(schema)
            .build(batch_stream)
            .map(|r| r.map_err(|e| Status::internal(e.to_string())));

        Ok(Response::new(Box::pin(flight_stream)))
    }

    async fn action_health_check(
        &self,
    ) -> Result<Response<BoxStream<'static, Result<arrow_flight::Result, Status>>>, Status> {
        let result = arrow_flight::Result { body: b"OK".to_vec().into() };
        Ok(Response::new(Box::pin(futures::stream::once(async { Ok(result) }))))
    }

    async fn action_get_partition_id(
        &self,
    ) -> Result<Response<BoxStream<'static, Result<arrow_flight::Result, Status>>>, Status> {
        let id = self.backend.read().await.partition_id();
        let result = arrow_flight::Result {
            body: id.to_string().into_bytes().into(),
        };
        Ok(Response::new(Box::pin(futures::stream::once(async { Ok(result) }))))
    }

    /// Load a partition from Parquet snapshot + WAL replay.
    ///
    /// Body JSON: `{ "partition_id": u32, "storage_path": "..." }`
    ///
    /// When compiled without `self-heal`, returns `Unimplemented`.
    async fn action_load_partition(
        &self,
        body: Vec<u8>,
    ) -> Result<Response<BoxStream<'static, Result<arrow_flight::Result, Status>>>, Status> {
        #[cfg(not(all(feature = "distributed", feature = "self-heal")))]
        {
            let _ = body;
            return Err(Status::unimplemented(
                "LoadPartition requires the `self-heal` feature",
            ));
        }

        #[cfg(all(feature = "distributed", feature = "self-heal"))]
        {
            #[derive(serde::Deserialize)]
            struct LoadReq {
                #[serde(default)]
                storage_path: String,
            }

            let req: LoadReq = serde_json::from_slice(&body)
                .map_err(|e| Status::invalid_argument(format!("LoadPartition body: {}", e)))?;

            let partition_id = self.backend.read().await.partition_id();

            tracing::info!(
                partition = partition_id,
                path = %req.storage_path,
                "LoadPartition action: restoring from snapshot"
            );

            let wal_config = WalConfig::from_env();
            let wal_cfg = if wal_config.enabled { Some(&wal_config) } else { None };

            let restored = snapshot_restore::restore_if_empty(
                Arc::clone(&self.backend),
                &req.storage_path,
                wal_cfg,
            )
            .await
            .map_err(|e| Status::internal(format!("LoadPartition restore failed: {}", e)))?;

            let msg = if restored {
                format!("partition {} restored from {}", partition_id, req.storage_path)
            } else {
                format!("partition {} already loaded or no snapshot at {}", partition_id, req.storage_path)
            };

            tracing::info!(%msg, "LoadPartition action complete");
            let result = arrow_flight::Result { body: msg.into_bytes().into() };
            Ok(Response::new(Box::pin(futures::stream::once(async { Ok(result) }))))
        }
    }

    /// Execute one BSP superstep for a distributed graph algorithm.
    ///
    /// Body JSON: [`SuperstepRequest`] — the algorithm name, superstep number,
    /// parameters, and incoming vertex messages.
    ///
    /// Response JSON: [`SuperstepResponse`] — outgoing messages and convergence flag.
    async fn action_run_superstep(
        &self,
        body: Vec<u8>,
    ) -> Result<Response<BoxStream<'static, Result<arrow_flight::Result, Status>>>, Status> {
        use crate::distributed::network::superstep::algorithms;

        let req: SuperstepRequest = serde_json::from_slice(&body)
            .map_err(|e| Status::invalid_argument(format!("RunSuperstep body: {}", e)))?;

        let partition_id = self.backend.read().await.partition_id();

        tracing::debug!(
            algo = %req.algo,
            superstep = req.superstep,
            partition = partition_id,
            "RunSuperstep"
        );

        // Gather local node IDs and edges for the algorithm
        let (local_node_ids, local_edges) = {
            let backend = self.backend.read().await;
            let node_ids: Vec<u64> = backend.all_nodes().map(|n| n.id()).collect();
            let edges: Vec<(u64, u64, f64)> = backend.all_edges()
                .filter_map(|e| {
                    backend.get_edge_endpoints(e.id()).map(|(s, d)| (s, d, 1.0_f64))
                })
                .collect();
            (node_ids, edges)
        };

        let resp = algorithms::dispatch(&req, &local_node_ids, &local_edges)
            .map_err(|e| Status::internal(format!("RunSuperstep algo error: {}", e)))?;

        let resp_json = serde_json::to_vec(&resp)
            .map_err(|e| Status::internal(format!("RunSuperstep serialize response: {}", e)))?;

        let result = arrow_flight::Result { body: resp_json.into() };
        Ok(Response::new(Box::pin(futures::stream::once(async { Ok(result) }))))
    }
}

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

#[cfg(all(test, feature = "distributed"))]
mod tests {
    use super::*;
    use crate::distributed::partition::PartitionMap;

    fn make_service(partition_id: u32) -> GraphFlightService {
        let map = Arc::new(std::sync::RwLock::new(PartitionMap::new(1)));
        let backend = PartitionedGraphBackend::new(partition_id, map);
        GraphFlightService::new(Arc::new(RwLock::new(backend)))
    }

    #[tokio::test]
    async fn test_health_check() {
        let svc = make_service(0);
        let action = Action { r#type: "HealthCheck".to_string(), body: vec![].into() };
        let mut stream = svc.do_action(Request::new(action)).await.unwrap().into_inner();
        let result = stream.next().await.unwrap().unwrap();
        assert_eq!(&result.body[..], b"OK");
    }

    #[tokio::test]
    async fn test_get_partition_id() {
        let svc = make_service(3);
        let action = Action { r#type: "GetPartitionId".to_string(), body: vec![].into() };
        let mut stream = svc.do_action(Request::new(action)).await.unwrap().into_inner();
        let result = stream.next().await.unwrap().unwrap();
        assert_eq!(String::from_utf8(result.body.to_vec()).unwrap(), "3");
    }

    #[tokio::test]
    async fn test_do_get_nodes_empty() {
        let svc = make_service(0);
        let ticket = Ticket { ticket: b"nodes".to_vec().into() };
        // Should return a valid (empty) flight stream without errors
        let resp = svc.do_get(Request::new(ticket)).await;
        assert!(resp.is_ok());
    }

    #[tokio::test]
    async fn test_do_get_edges_empty() {
        let svc = make_service(0);
        let ticket = Ticket { ticket: b"edges".to_vec().into() };
        let resp = svc.do_get(Request::new(ticket)).await;
        assert!(resp.is_ok());
    }

    #[tokio::test]
    async fn test_unknown_action_is_err() {
        let svc = make_service(0);
        let action = Action { r#type: "NotARealAction".to_string(), body: vec![].into() };
        let resp = svc.do_action(Request::new(action)).await;
        assert!(resp.is_err());
        assert_eq!(resp.err().unwrap().code(), tonic::Code::InvalidArgument);
    }

    #[tokio::test]
    async fn test_unknown_ticket_is_err() {
        let svc = make_service(0);
        let ticket = Ticket { ticket: b"something_else".to_vec().into() };
        let resp = svc.do_get(Request::new(ticket)).await;
        assert!(resp.is_err());
    }

    #[tokio::test]
    async fn test_list_actions() {
        let svc = make_service(0);
        let mut stream = svc.list_actions(Request::new(Empty {})).await.unwrap().into_inner();
        let first = stream.next().await.unwrap().unwrap();
        assert_eq!(first.r#type, "HealthCheck");
    }
}
