//! Strategies for selecting which pod should take over a failed partition.

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

// ── Types ─────────────────────────────────────────────────────────────────────

/// Snapshot of a running pod and its current load.
#[cfg(feature = "distributed")]
#[derive(Debug, Clone)]
pub struct PodInfo {
    /// Kubernetes pod name, e.g. `"ocg-partition-2"`.
    pub name:            String,
    /// Number of partitions currently assigned to this pod.
    pub partition_count: usize,
    /// gRPC endpoint for the pod's Flight server, e.g. `"http://ocg-partition-2:8815"`.
    pub endpoint:        String,
}

// ── Trait ─────────────────────────────────────────────────────────────────────

/// Strategy for choosing which pod recovers a failed partition.
#[cfg(feature = "distributed")]
pub trait PodSelector: Send + Sync {
    /// Select the best pod to host `partition_id` from the list of candidates.
    ///
    /// Returns `None` if `candidates` is empty.
    fn select(
        &self,
        partition_id: u32,
        candidates:   &[PodInfo],
        current_assignments: &HashMap<String, Vec<u32>>,
    ) -> Option<PodInfo>;
}

// ── Least-loaded selector ─────────────────────────────────────────────────────

/// Assigns the partition to whichever pod currently has the fewest partitions.
#[cfg(feature = "distributed")]
pub struct LeastLoadedSelector;

#[cfg(feature = "distributed")]
impl PodSelector for LeastLoadedSelector {
    fn select(
        &self,
        _partition_id: u32,
        candidates:    &[PodInfo],
        _assignments:  &HashMap<String, Vec<u32>>,
    ) -> Option<PodInfo> {
        candidates
            .iter()
            .min_by_key(|p| p.partition_count)
            .cloned()
    }
}

// ── StatefulSet ordinal selector ──────────────────────────────────────────────

/// Prefers the pod whose ordinal matches the partition ID.
///
/// For a StatefulSet named `ocg-partition` with replicas
/// `ocg-partition-0`, `ocg-partition-1`, …:
///
/// - partition 0 → prefers `ocg-partition-0`
/// - partition 5 → prefers `ocg-partition-5` (wraps if not available)
///
/// Falls back to `LeastLoadedSelector` when the preferred pod is absent.
#[cfg(feature = "distributed")]
pub struct StatefulOrdinalSelector {
    pub statefulset_name: String,
}

#[cfg(feature = "distributed")]
impl PodSelector for StatefulOrdinalSelector {
    fn select(
        &self,
        partition_id: u32,
        candidates:   &[PodInfo],
        assignments:  &HashMap<String, Vec<u32>>,
    ) -> Option<PodInfo> {
        if candidates.is_empty() {
            return None;
        }

        let preferred_name = format!("{}-{}", self.statefulset_name, partition_id);
        if let Some(pod) = candidates.iter().find(|p| p.name == preferred_name) {
            return Some(pod.clone());
        }

        // Fall back to least loaded
        LeastLoadedSelector.select(partition_id, candidates, assignments)
    }
}

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

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

    fn make_pods(names: &[&str]) -> Vec<PodInfo> {
        names
            .iter()
            .enumerate()
            .map(|(i, &name)| PodInfo {
                name:            name.to_string(),
                partition_count: i, // increasing load
                endpoint:        format!("http://{}:8815", name),
            })
            .collect()
    }

    #[test]
    fn test_least_loaded_picks_first() {
        let pods = make_pods(&["pod-0", "pod-1", "pod-2"]);
        let sel  = LeastLoadedSelector;
        let chosen = sel.select(0, &pods, &HashMap::new()).unwrap();
        assert_eq!(chosen.name, "pod-0"); // partition_count = 0
    }

    #[test]
    fn test_least_loaded_empty() {
        let sel = LeastLoadedSelector;
        assert!(sel.select(0, &[], &HashMap::new()).is_none());
    }

    #[test]
    fn test_stateful_ordinal_prefers_matching_pod() {
        let pods = make_pods(&["ocg-partition-0", "ocg-partition-1", "ocg-partition-2"]);
        let sel  = StatefulOrdinalSelector { statefulset_name: "ocg-partition".to_string() };
        let chosen = sel.select(2, &pods, &HashMap::new()).unwrap();
        assert_eq!(chosen.name, "ocg-partition-2");
    }

    #[test]
    fn test_stateful_ordinal_falls_back_to_least_loaded() {
        // Preferred pod (ocg-partition-5) is not in the candidate list
        let pods = make_pods(&["ocg-partition-0", "ocg-partition-1"]);
        let sel  = StatefulOrdinalSelector { statefulset_name: "ocg-partition".to_string() };
        let chosen = sel.select(5, &pods, &HashMap::new()).unwrap();
        assert_eq!(chosen.name, "ocg-partition-0"); // least loaded fallback
    }

    #[test]
    fn test_pod_info_fields() {
        let pod = PodInfo {
            name:            "ocg-partition-1".to_string(),
            partition_count: 2,
            endpoint:        "http://ocg-partition-1:8815".to_string(),
        };
        assert_eq!(pod.partition_count, 2);
        assert!(pod.endpoint.contains("8815"));
    }
}
