//! Dynamic partition rebalancer.
//!
//! Monitors partition metrics and triggers rebalancing when load becomes
//! imbalanced.  Currently implements the **detection** and **action planning**
//! stages; the actual data migration over Arrow Flight is wired in Phase 11+.
//!
//! ## Rebalancing triggers (configurable)
//! - `imbalance_threshold`:   max/avg node count ratio above which to trigger (default 1.2).
//! - `hot_partition_factor`:  query rate multiplier above which a partition is "hot" (default 2.0).
//!
//! ## Actions produced
//! - `Split`: divide a large partition into two smaller ones.
//! - `Merge`: combine two small partitions.
//! - `MigratePartition`: move a partition to a different pod.

#[cfg(feature = "distributed")]
use super::metrics::PartitionMetricsSnapshot;

// ── Config ────────────────────────────────────────────────────────────────────

/// Rebalancer configuration.
#[cfg(feature = "distributed")]
#[derive(Debug, Clone)]
pub struct RebalancerConfig {
    /// Trigger rebalancing when `max_nodes / avg_nodes > imbalance_threshold`.
    pub imbalance_threshold:  f64,
    /// A partition is "hot" when its query rate > `hot_partition_factor * avg_rate`.
    pub hot_partition_factor: f64,
    /// Minimum node count below which a partition is considered for merging.
    pub min_nodes_for_merge:  u64,
}

#[cfg(feature = "distributed")]
impl Default for RebalancerConfig {
    fn default() -> Self {
        Self {
            imbalance_threshold:  1.2,
            hot_partition_factor: 2.0,
            min_nodes_for_merge:  1_000,
        }
    }
}

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

/// Rebalancing action recommended by the planner.
#[cfg(feature = "distributed")]
#[derive(Debug, Clone, PartialEq)]
pub enum RebalanceAction {
    /// Split a large partition into two halves.
    Split {
        partition_id: u32,
        reason:       String,
    },
    /// Merge two small partitions.
    Merge {
        partition_a: u32,
        partition_b: u32,
        reason:      String,
    },
    /// Move a partition from one pod to another (hot partition relief).
    MigratePartition {
        partition_id:  u32,
        from_pod:      String,
        to_pod:        String,
        reason:        String,
    },
}

// ── Rebalancer ────────────────────────────────────────────────────────────────

/// Stateless rebalancing planner — converts metrics snapshots to actions.
#[cfg(feature = "distributed")]
pub struct Rebalancer {
    config: RebalancerConfig,
}

#[cfg(feature = "distributed")]
impl Rebalancer {
    pub fn new(config: RebalancerConfig) -> Self {
        Self { config }
    }

    /// Analyse partition metrics and return the set of recommended actions.
    pub fn plan(&self, snapshots: &[PartitionMetricsSnapshot]) -> Vec<RebalanceAction> {
        if snapshots.is_empty() {
            return vec![];
        }

        let mut actions = Vec::new();

        // ── Size imbalance ─────────────────────────────────────────────────
        let total_nodes: u64 = snapshots.iter().map(|s| s.node_count).sum();
        let avg_nodes = total_nodes as f64 / snapshots.len() as f64;
        let max_nodes = snapshots.iter().map(|s| s.node_count).max().unwrap_or(0);

        if avg_nodes > 0.0 && max_nodes as f64 / avg_nodes > self.config.imbalance_threshold {
            for snap in snapshots {
                if snap.node_count as f64 / avg_nodes > self.config.imbalance_threshold {
                    actions.push(RebalanceAction::Split {
                        partition_id: snap.partition_id,
                        reason: format!(
                            "node count {} exceeds {:.0}× average {:.0}",
                            snap.node_count, self.config.imbalance_threshold, avg_nodes
                        ),
                    });
                }
            }
        }

        // ── Merge candidates (very small partitions) ───────────────────────
        let small: Vec<u32> = snapshots
            .iter()
            .filter(|s| s.node_count < self.config.min_nodes_for_merge)
            .map(|s| s.partition_id)
            .collect();

        // Pair up small partitions for merging
        for pair in small.chunks(2) {
            if pair.len() == 2 {
                actions.push(RebalanceAction::Merge {
                    partition_a: pair[0],
                    partition_b: pair[1],
                    reason: format!(
                        "partitions {},{} are below minimum size {}",
                        pair[0], pair[1], self.config.min_nodes_for_merge
                    ),
                });
            }
        }

        // ── Hot partitions ─────────────────────────────────────────────────
        let total_queries: u64 = snapshots.iter().map(|s| s.query_count).sum();
        let avg_queries = total_queries as f64 / snapshots.len() as f64;

        if avg_queries > 0.0 {
            for snap in snapshots {
                if snap.query_count as f64 > avg_queries * self.config.hot_partition_factor {
                    // MigratePartition action: pod info not available at planning
                    // stage — the executor fills it in when applying the action.
                    actions.push(RebalanceAction::MigratePartition {
                        partition_id:  snap.partition_id,
                        from_pod:      "current".to_string(),
                        to_pod:        "least-loaded".to_string(),
                        reason: format!(
                            "query rate {} exceeds {:.0}× average {:.0}",
                            snap.query_count, self.config.hot_partition_factor, avg_queries
                        ),
                    });
                }
            }
        }

        actions
    }

    /// Return `true` if any rebalancing is needed.
    pub fn needs_rebalancing(&self, snapshots: &[PartitionMetricsSnapshot]) -> bool {
        !self.plan(snapshots).is_empty()
    }
}

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

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

    fn snap(partition_id: u32, node_count: u64, query_count: u64) -> PartitionMetricsSnapshot {
        PartitionMetricsSnapshot {
            partition_id,
            node_count,
            edge_count: node_count / 2,
            query_count,
            sampled_at: Instant::now(),
        }
    }

    #[test]
    fn test_balanced_no_action() {
        let rebalancer = Rebalancer::new(RebalancerConfig::default());
        let snaps = vec![
            snap(0, 10_000, 100),
            snap(1, 10_000, 100),
            snap(2, 10_000, 100),
        ];
        assert!(!rebalancer.needs_rebalancing(&snaps));
        assert!(rebalancer.plan(&snaps).is_empty());
    }

    #[test]
    fn test_imbalanced_triggers_split() {
        let rebalancer = Rebalancer::new(RebalancerConfig::default());
        // Partition 0 has 5x the nodes of partition 1 and 2
        let snaps = vec![
            snap(0, 100_000, 100),
            snap(1,  10_000, 100),
            snap(2,  10_000, 100),
        ];
        let actions = rebalancer.plan(&snaps);
        assert!(actions.iter().any(|a| matches!(a, RebalanceAction::Split { partition_id: 0, .. })));
    }

    #[test]
    fn test_small_partitions_trigger_merge() {
        let rebalancer = Rebalancer::new(RebalancerConfig::default());
        let snaps = vec![
            snap(0, 500, 10),   // below min_nodes_for_merge (1000)
            snap(1, 400, 10),   // below min_nodes_for_merge (1000)
            snap(2, 50_000, 10),
        ];
        let actions = rebalancer.plan(&snaps);
        assert!(actions.iter().any(|a| matches!(a, RebalanceAction::Merge { .. })));
    }

    #[test]
    fn test_hot_partition_triggers_migration() {
        let rebalancer = Rebalancer::new(RebalancerConfig::default());
        let snaps = vec![
            snap(0, 10_000, 1000), // hot: 1000 vs avg 370
            snap(1, 10_000, 10),
            snap(2, 10_000, 10),
            snap(3, 10_000, 80),
        ];
        let actions = rebalancer.plan(&snaps);
        assert!(actions.iter().any(|a| {
            matches!(a, RebalanceAction::MigratePartition { partition_id: 0, .. })
        }));
    }

    #[test]
    fn test_empty_snapshots() {
        let rebalancer = Rebalancer::new(RebalancerConfig::default());
        assert!(!rebalancer.needs_rebalancing(&[]));
        assert!(rebalancer.plan(&[]).is_empty());
    }

    #[test]
    fn test_custom_config() {
        let config = RebalancerConfig {
            imbalance_threshold:  1.5,
            hot_partition_factor: 3.0,
            min_nodes_for_merge:  500,
        };
        let rebalancer = Rebalancer::new(config);
        // With threshold 1.5, this should NOT trigger (max/avg ≈ 1.25)
        let snaps = vec![
            snap(0, 12_500, 10),
            snap(1, 10_000, 10),
            snap(2, 10_000, 10),
        ];
        assert!(!rebalancer.needs_rebalancing(&snaps));
    }
}
