//! WAL segment shipper: polls for sealed segments and ships them to peers.
//!
//! The shipper runs in a background thread, polling the WAL directory for
//! sealed segments (those with a newer segment after them) and sending each
//! to all configured peer shards via the [`Transport`] trait.

use std::path::PathBuf;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::thread::JoinHandle;
use std::time::Duration;

use crossbeam::channel::bounded;

use crate::replication::shard::ShardId;
use crate::replication::transport::{Transport, TransportError, WalSegmentPayload};
use crate::replication::{RegionId, WalSegmentId};
use crate::wal::format::batch::{HEADER_SIZE, decode_batch};

/// Configuration for the WAL shipper background thread.
pub struct ShipperConfig {
    /// Directory containing WAL segment files.
    pub wal_dir: PathBuf,
    /// This node's shard identity.
    pub shard_id: ShardId,
    /// Peer shards to ship segments to.
    pub peer_shards: Vec<ShardId>,
    /// How often to poll for new sealed segments.
    pub poll_interval: Duration,
}

impl Default for ShipperConfig {
    fn default() -> Self {
        Self {
            wal_dir: PathBuf::from("wal"),
            shard_id: ShardId::SINGLE,
            peer_shards: vec![],
            poll_interval: Duration::from_secs(2),
        }
    }
}

/// Handle to a running WAL shipper thread.
///
/// Call [`stop`](Self::stop) to signal shutdown and join the thread.
pub struct WalShipperHandle {
    shutdown_tx: crossbeam::channel::Sender<()>,
    thread: Option<JoinHandle<()>>,
}

impl WalShipperHandle {
    /// Signal the shipper to stop and join the background thread.
    pub fn stop(mut self) {
        let _ = self.shutdown_tx.send(());
        if let Some(handle) = self.thread.take() {
            let _ = handle.join();
        }
    }
}

/// Spawn a background thread that polls for sealed WAL segments and ships
/// them to all peer shards.
///
/// The shipper identifies sealed segments as those with a newer segment file
/// after them in sequence order. The most recent segment is skipped because
/// it may still be actively written.
///
/// # Panics
///
/// Panics if the OS fails to spawn the background thread.
pub fn spawn_shipper<T: Transport + ?Sized>(
    config: ShipperConfig,
    transport: Arc<T>,
) -> WalShipperHandle {
    let (shutdown_tx, shutdown_rx) = bounded::<()>(1);
    let thread = std::thread::Builder::new()
        .name("tidaldb-wal-shipper".into())
        .spawn(move || {
            let last_shipped = AtomicU64::new(0);
            loop {
                // Sleep for poll_interval, exit early on shutdown signal.
                match shutdown_rx.recv_timeout(config.poll_interval) {
                    Ok(()) | Err(crossbeam::channel::RecvTimeoutError::Disconnected) => {
                        tracing::debug!("wal shipper shutting down");
                        return;
                    }
                    Err(crossbeam::channel::RecvTimeoutError::Timeout) => {}
                }

                let segments = match crate::wal::segment::list_segments_for_shard(
                    &config.wal_dir,
                    config.shard_id,
                ) {
                    Ok(s) => s,
                    Err(e) => {
                        tracing::warn!(error = %e, "shipper: failed to list WAL segments");
                        continue;
                    }
                };

                // Only ship segments that have a newer segment after them
                // (the last segment may still be actively written).
                let sealed_count = if segments.len() > 1 {
                    segments.len() - 1
                } else {
                    0
                };
                let sealed = &segments[..sealed_count];
                let last = last_shipped.load(Ordering::Relaxed);

                for (seqno, path) in sealed {
                    if *seqno <= last {
                        continue;
                    }
                    let bytes = match std::fs::read(path) {
                        Ok(b) => b,
                        Err(e) => {
                            tracing::warn!(
                                error = %e,
                                path = %path.display(),
                                "shipper: failed to read segment"
                            );
                            continue;
                        }
                    };
                    let event_count = count_events_in_segment(&bytes);

                    for &peer in &config.peer_shards {
                        let payload = WalSegmentPayload {
                            id: WalSegmentId::new(RegionId::SINGLE, config.shard_id, *seqno),
                            bytes: bytes.clone(),
                            event_count,
                        };
                        match transport.send_segment(peer, payload) {
                            Ok(()) => {
                                tracing::debug!(
                                    seqno = *seqno,
                                    peer = %peer,
                                    "shipped WAL segment"
                                );
                            }
                            Err(TransportError::Closed) => {
                                tracing::warn!(
                                    peer = %peer,
                                    "shipper: transport closed, skipping peer"
                                );
                            }
                            Err(e) => {
                                tracing::warn!(
                                    error = %e,
                                    peer = %peer,
                                    "shipper: send failed"
                                );
                            }
                        }
                    }
                    last_shipped.store(*seqno, Ordering::Relaxed);
                }
            }
        })
        .expect("failed to spawn wal shipper thread");

    WalShipperHandle {
        shutdown_tx,
        thread: Some(thread),
    }
}

/// Count the total number of events in a WAL segment by parsing all batches.
fn count_events_in_segment(bytes: &[u8]) -> u64 {
    let mut offset = 0;
    let mut total = 0u64;
    while offset < bytes.len() {
        let remaining = &bytes[offset..];
        if remaining.len() < HEADER_SIZE {
            break;
        }
        match decode_batch(remaining) {
            Ok((header, _events)) => {
                total += u64::from(header.event_count);
                let batch_size = HEADER_SIZE + header.payload_len as usize;
                offset += batch_size;
            }
            Err(_) => break,
        }
    }
    total
}

#[cfg(test)]
#[allow(clippy::unwrap_used)]
mod tests {
    use super::*;
    use crate::wal::format::batch::{EventRecord, RECORD_TYPE_SIGNAL, encode_batch};

    #[test]
    fn count_events_empty() {
        assert_eq!(count_events_in_segment(&[]), 0);
    }

    #[test]
    fn count_events_single_batch() {
        let events = vec![EventRecord {
            entity_id: 1,
            signal_type: RECORD_TYPE_SIGNAL,
            weight: 1.0,
            timestamp_nanos: 100,
        }];
        let bytes = encode_batch(&events, 1, 1).unwrap();
        assert_eq!(count_events_in_segment(&bytes), 1);
    }

    #[test]
    fn count_events_multiple_batches() {
        let e1 = vec![EventRecord {
            entity_id: 1,
            signal_type: RECORD_TYPE_SIGNAL,
            weight: 1.0,
            timestamp_nanos: 100,
        }];
        let e2: Vec<EventRecord> = (0..5)
            .map(|i| EventRecord {
                entity_id: i,
                signal_type: RECORD_TYPE_SIGNAL,
                weight: 1.0,
                timestamp_nanos: 200,
            })
            .collect();
        let mut bytes = encode_batch(&e1, 1, 1).unwrap();
        bytes.extend(encode_batch(&e2, 2, 2).unwrap());
        assert_eq!(count_events_in_segment(&bytes), 6);
    }

    #[test]
    fn count_events_truncated_segment() {
        let events = vec![EventRecord {
            entity_id: 1,
            signal_type: RECORD_TYPE_SIGNAL,
            weight: 1.0,
            timestamp_nanos: 100,
        }];
        let mut bytes = encode_batch(&events, 1, 1).unwrap();
        // Append garbage that does not form a valid batch header.
        bytes.extend(&[0xFFu8; 10]);
        // Should still count the first valid batch.
        assert_eq!(count_events_in_segment(&bytes), 1);
    }

    #[test]
    fn shipper_config_default() {
        let config = ShipperConfig::default();
        assert_eq!(config.shard_id, ShardId::SINGLE);
        assert!(config.peer_shards.is_empty());
        assert_eq!(config.poll_interval, Duration::from_secs(2));
    }
}
