//! Recovery manager — coordinates partition migration after pod failure.
//!
//! ## Recovery flow
//!
//! 1. `FailureDetector` marks partition as `Failed` in etcd.
//! 2. `RecoveryManager::recover_partition` is called (manually or by a watch loop).
//! 3. Status set to `Recovering` in etcd.
//! 4. `PodSelector` chooses the best available pod.
//! 5. The chosen pod receives a gRPC `LoadPartition` call — it fetches the
//!    Parquet snapshot from object storage and loads it into memory.
//! 6. Pod registers itself as the new owner in etcd (`status = Healthy`).
//! 7. Future queries are routed to the new pod.

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

#[cfg(feature = "distributed")]
use tokio::sync::RwLock;

#[cfg(feature = "distributed")]
use crate::distributed::partition::location_service::{
    PartitionLocationService, PartitionStatus,
};
#[cfg(feature = "distributed")]
use crate::distributed::recovery::pod_selector::{PodInfo, PodSelector};
#[cfg(feature = "distributed")]
use crate::error::{CypherError, CypherResult, ExecutionError};

// ── Manager ───────────────────────────────────────────────────────────────────

/// Orchestrates partition recovery when pods fail.
#[cfg(feature = "distributed")]
pub struct RecoveryManager {
    location_svc: Arc<RwLock<PartitionLocationService>>,
    pod_selector: Box<dyn PodSelector>,
}

#[cfg(feature = "distributed")]
impl RecoveryManager {
    pub fn new(
        location_svc: Arc<RwLock<PartitionLocationService>>,
        pod_selector:  Box<dyn PodSelector>,
    ) -> Self {
        Self { location_svc, pod_selector }
    }

    // ── Public API ────────────────────────────────────────────────────────

    /// Recover all partitions currently marked as `Failed`.
    ///
    /// Returns the list of partition IDs that were successfully recovered.
    pub async fn recover_all_failed(
        &mut self,
        available_pods: &[PodInfo],
    ) -> CypherResult<Vec<u32>> {
        let failed = self.location_svc.read().await.failed_partitions().await;
        let mut recovered = Vec::new();
        for pid in failed {
            match self.recover_partition(pid, available_pods).await {
                Ok(())  => recovered.push(pid),
                Err(e)  => tracing::error!(
                    partition = pid,
                    error = %e,
                    "partition recovery failed"
                ),
            }
        }
        Ok(recovered)
    }

    /// Recover a single partition.
    ///
    /// Steps:
    /// 1. Mark as `Recovering`.
    /// 2. Select a new pod.
    /// 3. Notify the pod to load the partition (stub — real impl would call
    ///    the pod's `LoadPartition` Arrow Flight action).
    /// 4. Register the new pod as owner with `Healthy` status.
    pub async fn recover_partition(
        &mut self,
        partition_id:   u32,
        available_pods: &[PodInfo],
    ) -> CypherResult<()> {
        tracing::info!(partition = partition_id, "starting recovery");

        // Step 1: mark recovering
        self.location_svc
            .write()
            .await
            .update_partition_status(partition_id, PartitionStatus::Recovering)
            .await?;

        // Step 2: select new pod
        let current_assignments = self.build_assignment_map().await;
        let new_pod = self.pod_selector
            .select(partition_id, available_pods, &current_assignments)
            .ok_or_else(|| recovery_err("no available pods for recovery"))?;

        tracing::info!(
            partition = partition_id,
            pod = %new_pod.name,
            "selected recovery pod"
        );

        // Step 3: get storage path for the partition
        let storage_path = {
            let guard = self.location_svc.read().await;
            guard
                .get_partition_location(partition_id)
                .await
                .map(|loc| loc.storage_path.clone())
                .unwrap_or_default()
        };

        // Step 3: instruct new pod to load partition
        // In production this calls the pod's Arrow Flight "LoadPartition" action.
        // Here we call the stub which can be overridden in integration tests.
        self.load_partition_on_pod(&new_pod, partition_id, &storage_path).await?;

        // Step 4: register new owner
        self.location_svc
            .write()
            .await
            .register_partition(
                partition_id,
                new_pod.name.clone(),
                storage_path,
            )
            .await?;

        tracing::info!(
            partition = partition_id,
            pod = %new_pod.name,
            "recovery complete"
        );
        Ok(())
    }

    // ── Internal helpers ──────────────────────────────────────────────────

    async fn build_assignment_map(&self) -> HashMap<String, Vec<u32>> {
        let guard = self.location_svc.read().await;
        guard.assignment_map().await
    }

    /// Notify a pod to load a partition from object storage.
    ///
    /// Connects to the pod's Arrow Flight server and sends a `LoadPartition`
    /// action. The remote pod restores from Parquet snapshot + WAL replay.
    ///
    /// Retries up to 3 times with exponential backoff on transient failures.
    async fn load_partition_on_pod(
        &self,
        pod:          &PodInfo,
        partition_id: u32,
        storage_path: &str,
    ) -> CypherResult<()> {
        use crate::distributed::network::flight_client::GraphFlightClient;

        tracing::info!(
            pod      = %pod.name,
            endpoint = %pod.endpoint,
            partition = partition_id,
            path     = %storage_path,
            "load_partition_on_pod: connecting to pod"
        );

        let max_attempts = 3u32;
        let mut delay = std::time::Duration::from_millis(500);

        for attempt in 1..=max_attempts {
            match GraphFlightClient::connect(&pod.endpoint).await {
                Ok(mut client) => {
                    match client.load_partition(storage_path).await {
                        Ok(msg) => {
                            tracing::info!(
                                partition = partition_id,
                                pod = %pod.name,
                                response = %msg,
                                "load_partition_on_pod: success"
                            );
                            return Ok(());
                        }
                        Err(e) => {
                            if attempt == max_attempts {
                                return Err(recovery_err(format!(
                                    "LoadPartition on {} failed after {} attempts: {}",
                                    pod.name, max_attempts, e
                                )));
                            }
                            tracing::warn!(
                                attempt,
                                error = %e,
                                pod = %pod.name,
                                "LoadPartition failed, retrying"
                            );
                        }
                    }
                }
                Err(e) => {
                    if attempt == max_attempts {
                        return Err(recovery_err(format!(
                            "connect to {} failed after {} attempts: {}",
                            pod.endpoint, max_attempts, e
                        )));
                    }
                    tracing::warn!(
                        attempt,
                        error = %e,
                        endpoint = %pod.endpoint,
                        "Flight connect failed, retrying"
                    );
                }
            }
            tokio::time::sleep(delay).await;
            delay *= 2;
        }

        unreachable!()
    }
}

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

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

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

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

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

    #[test]
    fn test_recovery_err_helper() {
        let e = recovery_err("test error");
        assert!(matches!(e, CypherError::Execution(ExecutionError::Internal(_))));
    }

    #[tokio::test]
    #[ignore = "requires a running etcd server on localhost:2379"]
    async fn test_recover_partition_live() {
        use crate::distributed::partition::location_service::PartitionLocationService;

        let mut svc = PartitionLocationService::connect(
            vec!["http://127.0.0.1:2379".to_string()]
        ).await.unwrap();

        // Set up a failed partition
        svc.register_partition(77, "old-pod".to_string(), "s3://test/p77".to_string())
            .await.unwrap();
        svc.update_partition_status(77, PartitionStatus::Failed).await.unwrap();

        let svc = Arc::new(RwLock::new(svc));
        let mut mgr = RecoveryManager::new(
            Arc::clone(&svc),
            Box::new(LeastLoadedSelector),
        );

        let pods = make_pods(&["new-pod-0", "new-pod-1"]);
        mgr.recover_partition(77, &pods).await.unwrap();

        let owner = svc.read().await.get_partition_owner(77).await;
        assert_eq!(owner.as_deref(), Some("new-pod-0")); // least loaded
    }
}
