//! Pod startup sequence.
//!
//! ## Environment variables
//!
//! | Variable                   | Example                              | Description                         |
//! |----------------------------|--------------------------------------|-------------------------------------|
//! | `POD_NAME`                 | `ocg-partition-2`                    | Set by Kubernetes downward API      |
//! | `ETCD_ENDPOINTS`           | `http://etcd:2379`                   | Comma-separated etcd endpoints      |
//! | `S3_BUCKET`                | `ocg-graph-data`                     | Parquet snapshot storage bucket     |
//! | `S3_PREFIX`                | `graphs/my-graph`                    | Path prefix inside the bucket       |
//! | `PARTITION_IDS`            | `2,5`                                | Partitions this pod should own      |
//! | `FLIGHT_PORT`              | `8815`                               | Arrow Flight server port (default)  |
//! | `HEARTBEAT_INTERVAL`       | `5`                                  | Heartbeat seconds (default 5)       |
//! | `WAL_ENABLED`              | `true`                               | Enable WAL journaling (default on)  |
//! | `WAL_DIR`                  | `./wal`                              | Explicit WAL directory path         |
//! | `WAL_SYNC_MODE`            | `fsync`                              | `fsync` or `async`                  |
//! | `WAL_CHECKPOINT_THRESHOLD` | `10000`                              | Entries before auto-checkpoint      |
//! | `OCG_TLS_CERT`             | `/etc/ocg/tls/tls.crt`              | PEM certificate for TLS (opt-in)    |
//! | `OCG_TLS_KEY`              | `/etc/ocg/tls/tls.key`              | PEM private key for TLS             |
//! | `OCG_TLS_CA`               | `/etc/ocg/tls/ca.crt`               | CA cert for client auth (optional)  |
//!
//! When `PARTITION_IDS` is absent the pod derives its partition from the
//! StatefulSet ordinal embedded in `POD_NAME` (`ocg-partition-2` → 2).

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

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

#[cfg(feature = "distributed")]
use crate::distributed::graph::PartitionedGraphBackend;
#[cfg(feature = "distributed")]
use crate::distributed::health::HeartbeatService;
#[cfg(feature = "distributed")]
use crate::distributed::network::flight_server::GraphFlightService;
#[cfg(feature = "distributed")]
use crate::distributed::partition::location_service::PartitionLocationService;
#[cfg(feature = "distributed")]
use crate::distributed::partition::PartitionMap;
#[cfg(feature = "distributed")]
use crate::distributed::storage::wal::WalConfig;
#[cfg(feature = "distributed")]
use crate::error::{CypherError, CypherResult, ExecutionError};

// ── Configuration ─────────────────────────────────────────────────────────────

/// Pod runtime configuration, read from environment variables.
#[cfg(feature = "distributed")]
#[derive(Debug, Clone)]
pub struct PodConfig {
    pub pod_name:           String,
    pub etcd_endpoints:     Vec<String>,
    pub s3_bucket:          String,
    pub s3_prefix:          String,
    pub partition_ids:      Vec<u32>,
    pub flight_port:        u16,
    pub heartbeat_interval: u64,
}

#[cfg(feature = "distributed")]
impl PodConfig {
    /// Read configuration from environment variables.
    pub fn from_env() -> CypherResult<Self> {
        let pod_name = env_var("POD_NAME")?;

        // Parse ETCD_ENDPOINTS (comma-separated)
        let etcd_raw = env_var("ETCD_ENDPOINTS")
            .unwrap_or_else(|_| "http://etcd:2379".to_string());
        let etcd_endpoints: Vec<String> = etcd_raw
            .split(',')
            .map(|s| s.trim().to_string())
            .filter(|s| !s.is_empty())
            .collect();

        let s3_bucket = env_var("S3_BUCKET")
            .unwrap_or_else(|_| "ocg-graph-data".to_string());
        let s3_prefix = env_var("S3_PREFIX")
            .unwrap_or_else(|_| "graphs/default".to_string());

        // Parse PARTITION_IDS or derive from pod ordinal
        let partition_ids = if let Ok(ids_str) = env_var("PARTITION_IDS") {
            ids_str
                .split(',')
                .filter_map(|s| s.trim().parse::<u32>().ok())
                .collect()
        } else {
            // Derive from pod name: "ocg-partition-2" → [2]
            vec![ordinal_from_pod_name(&pod_name)?]
        };

        let flight_port = env_var("FLIGHT_PORT")
            .ok()
            .and_then(|s| s.parse().ok())
            .unwrap_or(8815u16);

        let heartbeat_interval = env_var("HEARTBEAT_INTERVAL")
            .ok()
            .and_then(|s| s.parse().ok())
            .unwrap_or(5u64);

        Ok(Self {
            pod_name,
            etcd_endpoints,
            s3_bucket,
            s3_prefix,
            partition_ids,
            flight_port,
            heartbeat_interval,
        })
    }

    /// Compute the S3 path for a given partition.
    pub fn partition_storage_path(&self, partition_id: u32) -> String {
        format!(
            "s3://{}/{}/partitions/{}",
            self.s3_bucket, self.s3_prefix, partition_id
        )
    }
}

// ── Startup ───────────────────────────────────────────────────────────────────

/// Orchestrates the full pod startup sequence.
#[cfg(feature = "distributed")]
pub struct PodStartup {
    config: PodConfig,
}

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

    /// Run the complete startup sequence:
    ///
    /// 1. Connect to etcd.
    /// 2. For each assigned partition: create an empty backend (real impl
    ///    would load from Parquet).
    /// 3. Register ownership in etcd.
    /// 4. Start the Arrow Flight server.
    /// 5. Start the heartbeat service.
    /// 6. Wait for a shutdown signal.
    pub async fn run(&self) -> CypherResult<()> {
        tracing::info!(
            pod         = %self.config.pod_name,
            partitions  = ?self.config.partition_ids,
            flight_port = self.config.flight_port,
            "pod starting up"
        );

        // 1. Connect to etcd (with exponential-backoff retry when self-heal is enabled)
        #[cfg(feature = "self-heal")]
        let etcd_client = crate::distributed::partition::etcd_retry::connect_with_retry(
            self.config.etcd_endpoints.clone(),
            std::env::var("ETCD_RETRY_MAX")
                .ok()
                .and_then(|s| s.parse().ok())
                .unwrap_or(10),
        )
        .await?;

        #[cfg(not(feature = "self-heal"))]
        let etcd_client = etcd_client::Client::connect(self.config.etcd_endpoints.clone(), None)
            .await
            .map_err(|e| startup_err(format!("etcd connect: {e}")))?;

        let location_svc =
            PartitionLocationService::from_client(etcd_client);
        let location_svc = Arc::new(RwLock::new(location_svc));

        // Build WAL configuration from environment variables.
        let wal_config = WalConfig::from_env();
        tracing::info!(
            wal_enabled = wal_config.enabled,
            wal_sync    = ?wal_config.sync_mode,
            "WAL configuration loaded"
        );

        // 2 & 3. For each partition: create backend + restore + attach WAL + register in etcd
        let partition_map = Arc::new(std::sync::RwLock::new(PartitionMap::new(
            self.config.partition_ids.len() as u32
        )));
        let mut backends: Vec<Arc<RwLock<PartitionedGraphBackend>>> = Vec::new();

        for &pid in &self.config.partition_ids {
            let storage_path = self.config.partition_storage_path(pid);

            // Create empty backend; snapshot restore (if enabled) runs below.
            let backend = PartitionedGraphBackend::new(pid, Arc::clone(&partition_map));
            let backend = Arc::new(RwLock::new(backend));

            // Attempt to restore from Parquet snapshot when self-heal is on.
            // Pass the WAL config so post-snapshot WAL entries are replayed.
            #[cfg(feature = "self-heal")]
            {
                match crate::distributed::storage::snapshot_restore::restore_if_empty(
                    Arc::clone(&backend),
                    &storage_path,
                    Some(&wal_config),
                )
                .await
                {
                    Ok(restored) => {
                        if restored {
                            tracing::info!(partition = pid, "partition restored from snapshot");
                        }
                    }
                    Err(e) => {
                        tracing::warn!(
                            partition = pid,
                            error     = %e,
                            "snapshot restore failed; continuing with empty partition"
                        );
                    }
                }
            }

            // Attach WAL to the backend so all future mutations are journaled.
            if wal_config.enabled {
                match crate::distributed::storage::wal::WriteAheadLog::open(&wal_config, &storage_path) {
                    Ok(wal) => {
                        backend.write().await.wal = Some(wal);
                        tracing::info!(partition = pid, "WAL attached to partition backend");
                    }
                    Err(e) => {
                        tracing::error!(
                            partition = pid,
                            error = %e,
                            "failed to open WAL; partition will run without journaling"
                        );
                    }
                }
            }

            backends.push(Arc::clone(&backend));

            location_svc
                .write()
                .await
                .register_partition(pid, self.config.pod_name.clone(), storage_path)
                .await?;

            tracing::info!(partition = pid, "partition registered");
        }

        // Keep a reference to all backends for graceful shutdown.
        let all_backends = backends.clone();

        // 4. Start Arrow Flight server
        // Use first partition backend (multi-partition pods can extend this)
        if let Some(backend) = backends.into_iter().next() {
            let service = GraphFlightService::new(backend).into_server();
            let addr    = format!("0.0.0.0:{}", self.config.flight_port)
                .parse()
                .map_err(|e| startup_err(format!("invalid address: {e}")))?;

            tracing::info!(%addr, "starting Arrow Flight server");

            // Spawn HTTP observability server (metrics / health / ready).
            let obs_state = crate::distributed::http::ObsState {
                metrics: crate::distributed::http::PartitionNodeMetrics::new(),
                ready: Arc::new(std::sync::atomic::AtomicBool::new(false)),
            };
            let http_addr = format!(
                "0.0.0.0:{}",
                std::env::var("HTTP_PORT")
                    .ok()
                    .and_then(|s| s.parse().ok())
                    .unwrap_or(9090u16)
            );
            let obs_clone = obs_state.clone();
            tokio::spawn(async move {
                if let Err(e) = crate::distributed::http::serve(&http_addr, obs_clone).await {
                    tracing::error!(error = %e, "HTTP observability server failed");
                }
            });

            // 5. Start heartbeat
            let mut hb = HeartbeatService::new(
                Arc::clone(&location_svc),
                self.config.partition_ids.clone(),
                self.config.heartbeat_interval,
            );
            hb.start()?;

            // Mark the pod as ready once Flight is about to serve traffic.
            obs_state.ready.store(true, std::sync::atomic::Ordering::Relaxed);

            // 6. Serve until shutdown signal — TLS if OCG_TLS_CERT / OCG_TLS_KEY are set
            let tls_cert = std::env::var("OCG_TLS_CERT").ok();
            let tls_key  = std::env::var("OCG_TLS_KEY").ok();

            match (tls_cert, tls_key) {
                (Some(cert_path), Some(key_path)) => {
                    let cert = tokio::fs::read(&cert_path).await
                        .map_err(|e| startup_err(format!("read TLS cert {cert_path}: {e}")))?;
                    let key = tokio::fs::read(&key_path).await
                        .map_err(|e| startup_err(format!("read TLS key {key_path}: {e}")))?;

                    let identity = tonic::transport::Identity::from_pem(cert, key);
                    let mut tls_config = tonic::transport::ServerTlsConfig::new()
                        .identity(identity);

                    // Optional: mutual TLS via CA cert
                    if let Ok(ca_path) = std::env::var("OCG_TLS_CA") {
                        let ca_cert = tokio::fs::read(&ca_path).await
                            .map_err(|e| startup_err(format!("read TLS CA {ca_path}: {e}")))?;
                        let ca = tonic::transport::Certificate::from_pem(ca_cert);
                        tls_config = tls_config.client_ca_root(ca);
                    }

                    tracing::info!(%addr, "starting Arrow Flight server (TLS)");
                    Server::builder()
                        .tls_config(tls_config)
                        .map_err(|e| startup_err(format!("TLS config error: {e}")))?
                        .add_service(service)
                        .serve_with_shutdown(addr, shutdown_signal())
                        .await
                        .map_err(|e| startup_err(format!("Flight server error: {e}")))?;
                }
                _ => {
                    tracing::info!(%addr, "starting Arrow Flight server (plain HTTP — set OCG_TLS_CERT/OCG_TLS_KEY for TLS)");
                    Server::builder()
                        .add_service(service)
                        .serve_with_shutdown(addr, shutdown_signal())
                        .await
                        .map_err(|e| startup_err(format!("Flight server error: {e}")))?;
                }
            }

            hb.stop();
        }

        // 7. Graceful shutdown: flush WAL and take final Parquet snapshot for each partition.
        for (i, backend_arc) in all_backends.iter().enumerate() {
            let pid = self.config.partition_ids[i];
            let storage_path = self.config.partition_storage_path(pid);

            let mut guard = backend_arc.write().await;

            // Flush WAL to ensure all buffered entries are on disk.
            if let Some(ref mut wal) = guard.wal {
                if let Err(e) = wal.checkpoint() {
                    tracing::error!(partition = pid, error = %e, "WAL checkpoint on shutdown failed");
                } else {
                    tracing::info!(partition = pid, "WAL checkpointed on shutdown");
                }
            }

            // Take a final Parquet snapshot so the next startup can skip WAL replay.
            let snapshot_dir = std::path::Path::new(&storage_path);
            match crate::distributed::storage::parquet_serde::save_partition_to_parquet(
                &*guard,
                snapshot_dir,
            ) {
                Ok(meta) => {
                    tracing::info!(
                        partition  = pid,
                        nodes     = meta.node_count,
                        edges     = meta.edge_count,
                        "final Parquet snapshot saved on shutdown"
                    );
                }
                Err(e) => {
                    tracing::error!(
                        partition = pid,
                        error    = %e,
                        "failed to save final Parquet snapshot on shutdown"
                    );
                }
            }
        }

        tracing::info!(pod = %self.config.pod_name, "pod shut down cleanly");
        Ok(())
    }
}

// ── Helpers ───────────────────────────────────────────────────────────────────

/// Extract the numeric ordinal from a Kubernetes StatefulSet pod name.
///
/// `"ocg-partition-3"` → `Ok(3)`
#[cfg(feature = "distributed")]
fn ordinal_from_pod_name(pod_name: &str) -> CypherResult<u32> {
    pod_name
        .rsplit('-')
        .next()
        .and_then(|s| s.parse::<u32>().ok())
        .ok_or_else(|| startup_err(format!(
            "cannot derive partition ordinal from POD_NAME '{pod_name}'; \
             set PARTITION_IDS explicitly"
        )))
}

#[cfg(feature = "distributed")]
fn env_var(name: &str) -> CypherResult<String> {
    std::env::var(name).map_err(|_| {
        startup_err(format!("required env var '{name}' not set"))
    })
}

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

/// Wait for SIGTERM or Ctrl-C.
#[cfg(feature = "distributed")]
async fn shutdown_signal() {
    use tokio::signal;
    let ctrl_c = async {
        signal::ctrl_c().await.expect("failed to install Ctrl+C handler");
    };

    #[cfg(unix)]
    let terminate = async {
        signal::unix::signal(signal::unix::SignalKind::terminate())
            .expect("failed to install SIGTERM handler")
            .recv()
            .await;
    };

    #[cfg(not(unix))]
    let terminate = std::future::pending::<()>();

    tokio::select! {
        _ = ctrl_c => {},
        _ = terminate => {},
    }
}

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

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

    #[test]
    fn test_ordinal_from_pod_name() {
        assert_eq!(ordinal_from_pod_name("ocg-partition-0").unwrap(), 0);
        assert_eq!(ordinal_from_pod_name("ocg-partition-3").unwrap(), 3);
        assert_eq!(ordinal_from_pod_name("ocg-partition-10").unwrap(), 10);
    }

    #[test]
    fn test_ordinal_bad_name() {
        assert!(ordinal_from_pod_name("ocg-partition-abc").is_err());
        assert!(ordinal_from_pod_name("nodigit").is_err());
    }

    #[test]
    fn test_partition_storage_path() {
        let cfg = PodConfig {
            pod_name:           "ocg-partition-2".to_string(),
            etcd_endpoints:     vec![],
            s3_bucket:          "my-bucket".to_string(),
            s3_prefix:          "graphs/g1".to_string(),
            partition_ids:      vec![2],
            flight_port:        8815,
            heartbeat_interval: 5,
        };
        assert_eq!(
            cfg.partition_storage_path(2),
            "s3://my-bucket/graphs/g1/partitions/2"
        );
    }

    #[test]
    fn test_pod_config_defaults() {
        // Just verify the struct can be constructed without panics
        let cfg = PodConfig {
            pod_name:           "test-pod-0".to_string(),
            etcd_endpoints:     vec!["http://localhost:2379".to_string()],
            s3_bucket:          "test-bucket".to_string(),
            s3_prefix:          "graphs/test".to_string(),
            partition_ids:      vec![0],
            flight_port:        8815,
            heartbeat_interval: 5,
        };
        assert_eq!(cfg.flight_port, 8815);
        assert_eq!(cfg.heartbeat_interval, 5);
    }
}
