armdb 0.7.0

use std::net::{SocketAddr, TcpStream};
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::thread::{self, JoinHandle};
use std::time::Duration;

use crate::error::{DbError, DbResult};
use crate::frame_reader::FrameReader;
use crate::shard::{Shard, Shards};
use crate::shutdown::ShutdownSignal;

use super::protocol::*;
use super::server::HEARTBEAT_INTERVAL_SECS;
use super::{ReplicationEntry, ReplicationRegistry};

const ACK_INTERVAL_CATCHUP: usize = 1; // Ack per batch during catch-up (C1: unblocks pipelined server)
const ACK_INTERVAL: usize = 1000; // Streaming: batched acks
const RECONNECT_BASE_MS: u64 = 1000;
const RECONNECT_MAX_MS: u64 = 30_000;

enum StreamFailure {
    Retryable(DbError),
    Terminal(DbError),
}

impl From<std::io::Error> for StreamFailure {
    fn from(error: std::io::Error) -> Self {
        if error.kind() == std::io::ErrorKind::InvalidData {
            Self::Terminal(DbError::Io(error))
        } else {
            Self::Retryable(DbError::Io(error))
        }
    }
}

impl From<DbError> for StreamFailure {
    fn from(error: DbError) -> Self {
        Self::Retryable(error)
    }
}

/// Replication client running on a follower node.
/// Connects to the leader and receives entries per-shard.
pub struct ReplicationClient {
    stop: ShutdownSignal,
    handles: Vec<JoinHandle<()>>,
}

/// Options for tuning the replication client (e.g. in tests).
pub struct ReplicationClientOptions {
    pub reconnect_base_ms: u64,
    pub reconnect_max_ms: u64,
    /// Follower-side liveness interval, mirrored from the leader's heartbeat
    /// interval (F4). The stream read timeout is `2×` this: if no frame (data or
    /// heartbeat) arrives within it, the leader is presumed dead and the client
    /// reconnects instead of blocking forever in `read_frame`. Should be ≥ the
    /// leader's configured heartbeat interval so ordinary heartbeats keep the
    /// connection alive.
    pub heartbeat_interval_secs: u64,
}

impl Default for ReplicationClientOptions {
    fn default() -> Self {
        Self {
            reconnect_base_ms: RECONNECT_BASE_MS,
            reconnect_max_ms: RECONNECT_MAX_MS,
            heartbeat_interval_secs: HEARTBEAT_INTERVAL_SECS,
        }
    }
}

impl ReplicationClient {
    /// Start the replication client.
    ///
    /// Spawns one thread per shard to connect to the leader and receive entries.
    pub fn start(
        leader_addr: SocketAddr,
        shards: Arc<Shards>,
        registry: Arc<ReplicationRegistry>,
        key_len: u16,
        signal: ShutdownSignal,
    ) -> DbResult<Self> {
        Self::start_with_options(
            leader_addr,
            shards,
            registry,
            key_len,
            signal,
            ReplicationClientOptions::default(),
        )
    }

    pub fn start_with_options(
        leader_addr: SocketAddr,
        shards: Arc<Shards>,
        registry: Arc<ReplicationRegistry>,
        key_len: u16,
        signal: ShutdownSignal,
        options: ReplicationClientOptions,
    ) -> DbResult<Self> {
        let mut handles = Vec::with_capacity(shards.len());
        let reconnect_base_ms = options.reconnect_base_ms;
        let reconnect_max_ms = options.reconnect_max_ms;
        let heartbeat_interval_secs = options.heartbeat_interval_secs;

        for shard_id in 0..shards.len() {
            let shards = shards.clone();
            let registry = registry.clone();
            let stop = signal.clone();

            let handle = thread::spawn(move || {
                run_shard_client(
                    leader_addr,
                    &shards,
                    shard_id,
                    &registry,
                    key_len,
                    &stop,
                    reconnect_base_ms,
                    reconnect_max_ms,
                    heartbeat_interval_secs,
                );
            });
            handles.push(handle);
        }

        Ok(Self {
            stop: signal,
            handles,
        })
    }

    pub fn stop(&self) {
        self.stop.shutdown();
    }
}

impl Drop for ReplicationClient {
    fn drop(&mut self) {
        self.stop.shutdown();
        for h in self.handles.drain(..) {
            let _ = h.join();
        }
    }
}

#[allow(clippy::too_many_arguments)]
fn run_shard_client(
    leader_addr: SocketAddr,
    shards: &[Shard],
    shard_id: usize,
    registry: &ReplicationRegistry,
    key_len: u16,
    stop: &ShutdownSignal,
    reconnect_base_ms: u64,
    reconnect_max_ms: u64,
    heartbeat_interval_secs: u64,
) {
    let shard = &shards[shard_id];
    let last_applied = Arc::new(AtomicU64::new(shard.durable_recovered_gsn()));
    let mut backoff_ms = reconnect_base_ms;

    loop {
        if stop.is_shutdown() {
            return;
        }

        match connect_and_stream(
            leader_addr,
            shards,
            shard_id,
            &last_applied,
            registry,
            key_len,
            stop.as_flag(),
            heartbeat_interval_secs,
        ) {
            Ok(()) => {
                tracing::info!(shard_id, "replication stream ended cleanly");
                return;
            }
            Err(StreamFailure::Retryable(e)) => {
                tracing::error!(shard_id, error = %e, backoff_ms, "replication error, reconnecting");
                if stop.wait_timeout(Duration::from_millis(backoff_ms)) {
                    return;
                }
                backoff_ms = (backoff_ms * 2).min(reconnect_max_ms);
            }
            Err(StreamFailure::Terminal(e)) => {
                tracing::error!(shard_id, error = %e, "terminal replication error");
                return;
            }
        }
    }
}

fn send_ack(writer: &mut TcpStream, shard_id: u8, last_gsn: u64) -> DbResult<()> {
    let ack = AckMessage { shard_id, last_gsn };
    write_frame(writer, &ack.encode())?;
    Ok(())
}

#[allow(clippy::too_many_arguments)]
fn connect_and_stream(
    leader_addr: SocketAddr,
    shards: &[Shard],
    shard_id: usize,
    last_applied: &AtomicU64,
    registry: &ReplicationRegistry,
    key_len: u16,
    stop: &std::sync::atomic::AtomicBool,
    heartbeat_interval_secs: u64,
) -> Result<(), StreamFailure> {
    let stream = TcpStream::connect_timeout(&leader_addr, Duration::from_secs(5))?;
    let _ = stream.set_nodelay(true);

    // F4: bound blocking reads so a silently dead leader (power-off with no RST)
    // cannot park this thread in `read_frame` forever. The leader heartbeats
    // every `heartbeat_interval_secs` during idle streaming, so 2× that window
    // with no frame at all means the connection is dead — surface it as an error
    // and let the caller back off and reconnect. This also lets the `stop` flag
    // be observed within one timeout period. Applies to every clone of the fd,
    // so the `reader` half below inherits it.
    stream.set_read_timeout(Some(Duration::from_secs(
        2 * heartbeat_interval_secs.max(1),
    )))?;

    let mut writer = stream.try_clone()?;
    let mut reader = stream;
    // Resumable reader for the apply loop (R2): the handshake below stays
    // on the free `read_frame` — read_exact never reads ahead, so nothing
    // can be lost. A frame torn by the read timeout must not silently
    // desynchronize the stream if this loop ever continues after a timeout.
    let mut frames = FrameReader::<MessageType>::new();

    // Send SyncRequest — ask for the entry *after* last_gsn so the leader
    // does not resend the already-applied entry (C12).
    let req = SyncRequest {
        protocol_version: VAR_PROTOCOL_VERSION,
        shard_id: shard_id as u8,
        from_gsn: last_applied.load(Ordering::Relaxed).saturating_add(1),
        key_len,
    };
    write_frame(&mut writer, &req.encode())?;

    // Receive ShardInfo
    let info_frame = read_frame(&mut reader)?;
    if info_frame.msg_type == MessageType::Error {
        return Err(classify_remote_error(&info_frame.payload));
    }
    if info_frame.msg_type != MessageType::ShardInfo {
        return Err(StreamFailure::Terminal(DbError::Replication(format!(
            "expected ShardInfo, got {:?}",
            info_frame.msg_type
        ))));
    }
    let info = ShardInfo::decode(&info_frame.payload).map_err(|error| {
        StreamFailure::Terminal(DbError::Replication(format!("invalid ShardInfo: {error}")))
    })?;
    if info.protocol_version != VAR_PROTOCOL_VERSION {
        return Err(StreamFailure::Terminal(DbError::Replication(format!(
            "variable replication protocol mismatch: leader {}, follower {}",
            info.protocol_version, VAR_PROTOCOL_VERSION
        ))));
    }
    if info.shard_count as usize != shards.len() {
        return Err(StreamFailure::Terminal(DbError::ShardCountMismatch {
            leader: info.shard_count as usize,
            follower: shards.len(),
        }));
    }

    tracing::info!(
        shard_id,
        from_gsn = last_applied.load(Ordering::Relaxed),
        "connected to leader"
    );

    let shard = &shards[shard_id];
    let mut entries_since_ack = 0usize;
    // Track catch-up phase: true until the server sends CaughtUp (C1).
    let mut in_catchup = true;

    loop {
        if stop.load(Ordering::Relaxed) {
            return Ok(());
        }

        let frame = frames.read_frame(&mut reader)?;

        match frame.msg_type {
            MessageType::EntryBatch => {
                let batch = EntryBatch::decode(&frame.payload)?;
                // C15: reject batches whose shard_id doesn't match this stream.
                if batch.shard_id != shard_id as u8 {
                    return Err(StreamFailure::Terminal(DbError::Replication(
                        "EntryBatch shard_id mismatch".into(),
                    )));
                }

                for wire_entry in &batch.entries {
                    let repl_entry = ReplicationEntry {
                        data: wire_entry.data.clone(),
                        key_len: wire_entry.key_len,
                    };
                    registry.apply_streaming(shard, &repl_entry, last_applied)?;

                    entries_since_ack += 1;
                }

                // Send Ack at phase-appropriate interval.
                // During catch-up: ack every batch (ACK_INTERVAL_CATCHUP=1) so the
                // server's catch-up loop is never stalled waiting for an ack (C1).
                // During streaming: ack every ACK_INTERVAL entries (batched).
                let ack_threshold = if in_catchup {
                    ACK_INTERVAL_CATCHUP
                } else {
                    ACK_INTERVAL
                };
                if entries_since_ack >= ack_threshold {
                    send_ack(
                        &mut writer,
                        shard_id as u8,
                        last_applied.load(Ordering::Relaxed),
                    )?;
                    entries_since_ack = 0;
                }
            }
            MessageType::CaughtUp => {
                let caught_up = CaughtUp::decode(&frame.payload)?;
                tracing::info!(shard_id, leader_gsn = caught_up.leader_gsn, "caught up");
                in_catchup = false;
                // Continue — streaming will follow
            }
            MessageType::Heartbeat => {
                send_ack(
                    &mut writer,
                    shard_id as u8,
                    last_applied.load(Ordering::Relaxed),
                )?;
                entries_since_ack = 0;
            }
            MessageType::Error => {
                return Err(classify_remote_error(&frame.payload));
            }
            other => {
                return Err(StreamFailure::Terminal(DbError::Replication(format!(
                    "unexpected message: {other:?}"
                ))));
            }
        }
    }
}

fn classify_remote_error(payload: &[u8]) -> StreamFailure {
    let error = match decode_error(payload) {
        Ok(error) => error,
        Err(error) => {
            return StreamFailure::Terminal(DbError::Replication(format!(
                "malformed replication error: {error}"
            )));
        }
    };

    let db_error = DbError::Replication(error.message);
    match error.code {
        ReplicationErrorCode::RetryableIo => StreamFailure::Retryable(db_error),
        ReplicationErrorCode::ResourceLimit
        | ReplicationErrorCode::CorruptedLog
        | ReplicationErrorCode::ProtocolMismatch
        | ReplicationErrorCode::InvalidRequest => StreamFailure::Terminal(db_error),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::{self, Cursor};

    fn assert_terminal(failure: StreamFailure) {
        assert!(matches!(failure, StreamFailure::Terminal(_)));
    }

    fn assert_retryable(failure: StreamFailure) {
        assert!(matches!(failure, StreamFailure::Retryable(_)));
    }

    #[test]
    fn malformed_protocol_decode_errors_are_terminal() {
        let entry_batch = EntryBatch::decode(&[])
            .err()
            .expect("empty EntryBatch must fail");
        assert_terminal(StreamFailure::from(entry_batch));

        let caught_up = CaughtUp::decode(&[])
            .err()
            .expect("empty CaughtUp must fail");
        assert_terminal(StreamFailure::from(caught_up));

        let typed_error = decode_error(&[]).expect_err("empty typed error must fail");
        assert_terminal(StreamFailure::from(typed_error));
        assert_terminal(classify_remote_error(&[]));

        let unknown_frame = read_frame(&mut Cursor::new([0])).expect_err("unknown frame must fail");
        assert_terminal(StreamFailure::from(unknown_frame));

        let mut oversized = Vec::from([MessageType::EntryBatch as u8]);
        oversized.extend_from_slice(&((MAX_FRAME_SIZE + 1) as u32).to_le_bytes());
        let oversized_frame =
            read_frame(&mut Cursor::new(oversized)).expect_err("oversized frame must fail");
        assert_terminal(StreamFailure::from(oversized_frame));
    }

    #[test]
    fn transport_io_is_retryable() {
        assert_retryable(StreamFailure::from(io::Error::new(
            io::ErrorKind::TimedOut,
            "socket timeout",
        )));
    }

    #[test]
    fn remote_retryable_io_is_retryable() {
        let frame = encode_error(ReplicationErrorCode::RetryableIo, "temporary read failure");
        assert_retryable(classify_remote_error(&frame.payload));
    }

    #[test]
    fn remote_non_retryable_codes_are_terminal() {
        for code in [
            ReplicationErrorCode::ResourceLimit,
            ReplicationErrorCode::CorruptedLog,
            ReplicationErrorCode::ProtocolMismatch,
            ReplicationErrorCode::InvalidRequest,
        ] {
            let frame = encode_error(code, "fatal replication failure");
            assert_terminal(classify_remote_error(&frame.payload));
        }
    }
}