//! BSP (Bulk Synchronous Parallel) superstep framework over Arrow Flight.
//!
//! ## Protocol
//!
//! The coordinator drives execution by sending `RunSuperstep` DoActions to
//! each partition pod via Arrow Flight.  Each round proceeds as follows:
//!
//! 1. Coordinator serialises a [`SuperstepRequest`] to JSON and calls
//!    `DoAction("RunSuperstep", body)` on every pod.
//! 2. Each pod deserialises the request, runs its local algorithm step
//!    (see [`BspAlgorithm`]), and returns a [`SuperstepResponse`] with
//!    outgoing messages and a convergence flag.
//! 3. Coordinator collects all responses, routes messages to their target
//!    partition, and begins the next superstep — or stops if all pods
//!    report convergence.
//!
//! Messages are typed as `f64` scalars (sufficient for PageRank, SSSP,
//! Betweenness accumulators).  Richer message types can be added later.

#[cfg(feature = "distributed")]
use std::collections::HashMap;

#[cfg(feature = "distributed")]
use serde::{Deserialize, Serialize};

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

// ── Message types ─────────────────────────────────────────────────────────────

/// A vertex-to-vertex message in the BSP model.
///
/// `dst_vertex` is the *global* stable vertex ID in the receiving partition.
/// `value` is the message payload (a single f64 scalar is sufficient for
/// PageRank, SSSP, and betweenness accumulators).
#[cfg(feature = "distributed")]
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct VertexMessage {
    pub dst_vertex: u64,
    pub value: f64,
}

/// Request sent to a pod for one BSP superstep.
#[cfg(feature = "distributed")]
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SuperstepRequest {
    /// Name of the algorithm running (e.g. `"pagerank"`, `"sssp"`, `"bfs"`).
    pub algo: String,

    /// Current superstep number (0-indexed).
    pub superstep: u32,

    /// Algorithm-specific parameters (convergence threshold, damping factor, …).
    pub params: HashMap<String, f64>,

    /// Incoming messages for vertices in *this* partition.
    ///
    /// Keyed by the destination vertex's global stable ID.
    pub messages: HashMap<u64, Vec<f64>>,
}

/// Response from a pod after completing one BSP superstep.
#[cfg(feature = "distributed")]
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SuperstepResponse {
    /// The superstep that was just executed.
    pub superstep: u32,

    /// Outgoing messages to vertices in *other* partitions.
    ///
    /// Destination vertex IDs are global stable IDs; the coordinator routes
    /// each message to the correct partition.
    pub outgoing: Vec<VertexMessage>,

    /// `true` when this partition has no more work to do (local convergence).
    ///
    /// The overall computation halts only when *all* partitions report `true`.
    pub converged: bool,

    /// Optional per-vertex result values to include in the final aggregation.
    ///
    /// Populated only on the last superstep (when `converged == true`).
    /// Maps global vertex ID → computed value (rank, distance, etc.).
    pub local_results: HashMap<u64, f64>,
}

// ── BSP algorithm trait ───────────────────────────────────────────────────────

/// A local algorithm that can participate in BSP superstep execution.
///
/// Implementations process the incoming vertex messages for their partition,
/// update local state, and produce outgoing messages for vertices in other
/// partitions.
#[cfg(feature = "distributed")]
pub trait BspAlgorithm: Send + Sync {
    /// Run one BSP superstep.
    ///
    /// # Parameters
    /// - `graph`      — read-only reference to the local partition's graph data
    /// - `request`    — incoming request with messages and algorithm parameters
    ///
    /// # Returns
    /// A [`SuperstepResponse`] containing outgoing messages and convergence state.
    fn run_superstep(
        &mut self,
        request: &SuperstepRequest,
        local_node_ids: &[u64],
        local_edges: &[(u64, u64, f64)], // (src, dst, weight) — cross-partition edges included
    ) -> CypherResult<SuperstepResponse>;
}

// ── Coordinator ───────────────────────────────────────────────────────────────

/// Drives BSP computation across all partition pods.
///
/// The coordinator connects to each partition's Flight server, sends superstep
/// requests, collects responses, and routes inter-partition messages until all
/// partitions converge.
#[cfg(feature = "distributed")]
pub struct BspCoordinator {
    /// Map from partition_id → Flight endpoint URL.
    partition_endpoints: HashMap<u32, String>,

    /// Maximum number of supersteps before declaring non-convergence.
    pub max_supersteps: u32,
}

#[cfg(feature = "distributed")]
impl BspCoordinator {
    /// Create a coordinator with the given partition → endpoint mapping.
    pub fn new(partition_endpoints: HashMap<u32, String>) -> Self {
        Self {
            partition_endpoints,
            max_supersteps: 100,
        }
    }

    /// Run a BSP algorithm to completion across all partitions.
    ///
    /// `algo_name` — name registered in the partition's `RunSuperstep` handler.
    /// `params`    — algorithm parameters (e.g. `{"damping": 0.85}`).
    ///
    /// Returns a map from global vertex ID → final computed value.
    pub async fn run(
        &self,
        algo: &str,
        params: HashMap<String, f64>,
    ) -> CypherResult<HashMap<u64, f64>> {
        use super::flight_client::GraphFlightClient;

        // Connect to all partitions
        let mut clients: HashMap<u32, GraphFlightClient> = HashMap::new();
        for (&pid, endpoint) in &self.partition_endpoints {
            let client = GraphFlightClient::connect(endpoint.clone()).await
                .map_err(|e| internal(format!("connect to partition {pid}: {e}")))?;
            clients.insert(pid, client);
        }

        // Messages in flight: partition_id → (vertex_id → [values])
        let mut pending: HashMap<u32, HashMap<u64, Vec<f64>>> = HashMap::new();

        let mut all_results: HashMap<u64, f64> = HashMap::new();

        for superstep in 0..self.max_supersteps {
            let mut all_converged = true;
            let mut next_pending: HashMap<u32, HashMap<u64, Vec<f64>>> = HashMap::new();

            // Issue superstep to all partitions in parallel
            // (For simplicity we do this sequentially here; a production
            //  implementation would use tokio::join_all)
            for (&pid, client) in &mut clients {
                let inbox = pending.remove(&pid).unwrap_or_default();
                let req = SuperstepRequest {
                    algo: algo.to_string(),
                    superstep,
                    params: params.clone(),
                    messages: inbox,
                };

                let resp = client.run_superstep(&req).await
                    .map_err(|e| internal(format!("superstep {superstep} on partition {pid}: {e}")))?;

                if !resp.converged {
                    all_converged = false;
                }

                // Route outgoing messages to the correct partition
                for msg in resp.outgoing {
                    // Simple round-robin routing: partition = vertex_id % num_partitions
                    // In a real system, use the PartitionMap for accurate routing
                    let target = (msg.dst_vertex as u32) % (self.partition_endpoints.len() as u32);
                    next_pending
                        .entry(target)
                        .or_default()
                        .entry(msg.dst_vertex)
                        .or_default()
                        .push(msg.value);
                }

                // Collect local results (populated on last superstep)
                all_results.extend(resp.local_results);
            }

            pending = next_pending;

            if all_converged && pending.values().all(|m| m.is_empty()) {
                tracing::info!(superstep, algo, "BSP converged");
                break;
            }

            if superstep == self.max_supersteps - 1 {
                tracing::warn!(superstep, algo, "BSP did not converge within max supersteps");
            }
        }

        Ok(all_results)
    }
}

// ── Error helper ──────────────────────────────────────────────────────────────

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

// ── Algorithm dispatch ────────────────────────────────────────────────────────

/// Built-in BSP algorithm implementations.
///
/// Each algorithm implements the [`BspAlgorithm`] pattern as a stateless
/// function for simplicity (state is carried in the message accumulation).
#[cfg(feature = "distributed")]
pub mod algorithms {
    use super::*;

    /// Dispatch a superstep request to the appropriate built-in algorithm.
    ///
    /// Supported algorithm names (case-insensitive):
    /// - `"pagerank"` — iterative PageRank
    /// - `"bfs"` — BFS level propagation
    /// - `"sssp"` — single-source shortest paths (Bellman-Ford style)
    pub fn dispatch(
        req: &SuperstepRequest,
        local_node_ids: &[u64],
        local_edges: &[(u64, u64, f64)],
    ) -> CypherResult<SuperstepResponse> {
        match req.algo.to_lowercase().as_str() {
            "pagerank" => pagerank_superstep(req, local_node_ids, local_edges),
            "bfs"      => bfs_superstep(req, local_node_ids, local_edges),
            "sssp"     => sssp_superstep(req, local_node_ids, local_edges),
            other => Err(CypherError::Execution(
                crate::error::ExecutionError::Internal(
                    format!("Unknown BSP algorithm: {other}")
                ),
            )),
        }
    }

    /// PageRank BSP superstep.
    ///
    /// Each vertex accumulates incoming rank contributions, updates its own
    /// rank, and distributes 1/out_degree to each out-neighbor.
    ///
    /// Parameters:
    /// - `damping` — damping factor (default 0.85)
    /// - `n_total` — total number of vertices in the graph (for the dangling-node correction)
    fn pagerank_superstep(
        req: &SuperstepRequest,
        local_node_ids: &[u64],
        local_edges: &[(u64, u64, f64)],
    ) -> CypherResult<SuperstepResponse> {
        let damping = *req.params.get("damping").unwrap_or(&0.85_f64);
        let n_total = *req.params.get("n_total").unwrap_or(&1.0_f64) as usize;
        let n_total = n_total.max(1);

        // Compute out-degree map
        let mut out_degree: HashMap<u64, usize> = HashMap::new();
        for &nid in local_node_ids {
            out_degree.insert(nid, 0);
        }
        for &(src, _, _) in local_edges {
            *out_degree.entry(src).or_insert(0) += 1;
        }

        // Accumulate incoming rank for each vertex
        let mut rank: HashMap<u64, f64> = HashMap::new();
        for (&vid, values) in &req.messages {
            let sum: f64 = values.iter().sum();
            rank.insert(vid, sum);
        }

        // For vertices with no incoming message this superstep, use teleportation
        for &nid in local_node_ids {
            rank.entry(nid).or_insert(0.0);
        }

        // Apply damping and teleportation
        let teleport = (1.0 - damping) / (n_total as f64);
        for v in rank.values_mut() {
            *v = damping * *v + teleport;
        }

        // Send rank / out_degree to each out-neighbor
        let mut outgoing: Vec<VertexMessage> = Vec::new();
        for &(src, dst, _weight) in local_edges {
            if let Some(&r) = rank.get(&src) {
                let out_deg = *out_degree.get(&src).unwrap_or(&1);
                outgoing.push(VertexMessage {
                    dst_vertex: dst,
                    value: r / (out_deg as f64),
                });
            }
        }

        // PageRank never formally converges from the algorithm's perspective;
        // convergence is detected by the coordinator comparing value changes.
        // We signal convergence after a fixed number of iterations (via params).
        let max_iters = *req.params.get("max_iters").unwrap_or(&20.0) as u32;
        let converged = req.superstep >= max_iters;

        let local_results = if converged { rank.clone() } else { HashMap::new() };

        Ok(SuperstepResponse {
            superstep: req.superstep,
            outgoing,
            converged,
            local_results,
        })
    }

    /// BFS level-propagation superstep.
    ///
    /// Each vertex that has been reached (level ≥ 0) propagates `level + 1`
    /// to its unvisited neighbors.
    ///
    /// Parameters:
    /// - `source` — source vertex ID (used only in superstep 0)
    fn bfs_superstep(
        req: &SuperstepRequest,
        local_node_ids: &[u64],
        local_edges: &[(u64, u64, f64)],
    ) -> CypherResult<SuperstepResponse> {
        let source = *req.params.get("source").unwrap_or(&0.0) as u64;

        // Our "level" state is carried in the messages.
        // level[v] = minimum incoming level value (i.e., shortest path length)
        let mut level: HashMap<u64, f64> = HashMap::new();

        // Superstep 0: seed the source vertex
        if req.superstep == 0 {
            if local_node_ids.contains(&source) {
                level.insert(source, 0.0);
            }
        }

        // Apply incoming messages (take minimum level)
        for (&vid, values) in &req.messages {
            let min_val = values.iter().cloned().fold(f64::INFINITY, f64::min);
            let entry = level.entry(vid).or_insert(f64::INFINITY);
            if min_val < *entry {
                *entry = min_val;
            }
        }

        // Propagate to out-neighbors
        let mut outgoing: Vec<VertexMessage> = Vec::new();
        for &(src, dst, _) in local_edges {
            if let Some(&lv) = level.get(&src) {
                outgoing.push(VertexMessage { dst_vertex: dst, value: lv + 1.0 });
            }
        }

        // Converge when no new messages are produced
        let converged = outgoing.is_empty() && req.messages.is_empty() && req.superstep > 0;

        let local_results = if converged {
            level.iter().filter(|(_, &v)| v < f64::INFINITY).map(|(&k, &v)| (k, v)).collect()
        } else {
            level.iter().filter(|(_, &v)| v < f64::INFINITY).map(|(&k, &v)| (k, v)).collect()
        };

        Ok(SuperstepResponse {
            superstep: req.superstep,
            outgoing,
            converged,
            local_results,
        })
    }

    /// SSSP (Bellman-Ford style) superstep.
    ///
    /// Each vertex propagates `dist + weight` to each out-neighbor.
    /// Convergence when no updates occur.
    ///
    /// Parameters:
    /// - `source` — source vertex ID
    fn sssp_superstep(
        req: &SuperstepRequest,
        local_node_ids: &[u64],
        local_edges: &[(u64, u64, f64)],
    ) -> CypherResult<SuperstepResponse> {
        let source = *req.params.get("source").unwrap_or(&0.0) as u64;

        let mut dist: HashMap<u64, f64> = HashMap::new();

        // Superstep 0: initialise source
        if req.superstep == 0 && local_node_ids.contains(&source) {
            dist.insert(source, 0.0);
        }

        // Apply incoming messages (take minimum distance)
        for (&vid, values) in &req.messages {
            let min_val = values.iter().cloned().fold(f64::INFINITY, f64::min);
            let entry = dist.entry(vid).or_insert(f64::INFINITY);
            if min_val < *entry {
                *entry = min_val;
            }
        }

        // Relax edges
        let mut outgoing: Vec<VertexMessage> = Vec::new();
        for &(src, dst, w) in local_edges {
            if let Some(&d) = dist.get(&src) {
                let new_dist = d + w;
                outgoing.push(VertexMessage { dst_vertex: dst, value: new_dist });
            }
        }

        let converged = outgoing.is_empty() && req.messages.is_empty() && req.superstep > 0;

        let local_results = dist.iter().filter(|(_, &v)| v < f64::INFINITY).map(|(&k, &v)| (k, v)).collect();

        Ok(SuperstepResponse {
            superstep: req.superstep,
            outgoing,
            converged,
            local_results,
        })
    }
}

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

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

    fn make_request(superstep: u32, messages: HashMap<u64, Vec<f64>>) -> SuperstepRequest {
        SuperstepRequest {
            algo: "pagerank".to_string(),
            superstep,
            params: [("damping".to_string(), 0.85_f64)].into(),
            messages,
        }
    }

    #[test]
    fn test_superstep_request_round_trip() {
        let msgs: HashMap<u64, Vec<f64>> = [(1u64, vec![0.5_f64, 0.25_f64])].into();
        let req = make_request(3, msgs);
        let json = serde_json::to_string(&req).unwrap();
        let decoded: SuperstepRequest = serde_json::from_str(&json).unwrap();
        assert_eq!(decoded.superstep, 3);
        assert_eq!(decoded.algo, "pagerank");
        assert_eq!(decoded.messages[&1], vec![0.5, 0.25]);
        assert!((decoded.params["damping"] - 0.85).abs() < 1e-9);
    }

    #[test]
    fn test_superstep_response_round_trip() {
        let resp = SuperstepResponse {
            superstep: 2,
            outgoing: vec![VertexMessage { dst_vertex: 42, value: 0.33 }],
            converged: false,
            local_results: [(7u64, 0.11_f64)].into(),
        };
        let json = serde_json::to_string(&resp).unwrap();
        let decoded: SuperstepResponse = serde_json::from_str(&json).unwrap();
        assert_eq!(decoded.superstep, 2);
        assert!(!decoded.converged);
        assert_eq!(decoded.outgoing[0].dst_vertex, 42);
        assert!((decoded.local_results[&7] - 0.11).abs() < 1e-9);
    }

    #[test]
    fn test_coordinator_new() {
        let endpoints: HashMap<u32, String> = [
            (0, "http://localhost:8815".to_string()),
            (1, "http://localhost:8816".to_string()),
        ].into();
        let coord = BspCoordinator::new(endpoints);
        assert_eq!(coord.max_supersteps, 100);
        assert_eq!(coord.partition_endpoints.len(), 2);
    }
}
