//! Segment receiver: consumes WAL segments from the transport and applies
//! them to the local signal ledger.
//!
//! The receiver runs in a background thread, blocking on
//! [`Transport::recv_segment`] and replaying each batch into the shared
//! [`SignalLedger`] via `apply_wal_event`. Idempotent replay is ensured by
//! checking the per-shard high-water-mark in [`ReplicationState`].

use std::sync::Arc;
use std::thread::JoinHandle;

use crate::replication::shard::ShardId;
use crate::replication::state::ReplicationState;
use crate::replication::transport::Transport;
use crate::schema::{EntityId, Timestamp};
use crate::signals::{SignalLedger, SignalTypeId};
use crate::wal::format::batch::{HEADER_SIZE, decode_batch};

/// Handle to a running segment receiver thread.
///
/// Call [`join`](Self::join) to block until the thread exits (triggered by
/// the transport returning `None` from `recv_segment`).
pub struct SegmentReceiverHandle {
    thread: Option<JoinHandle<()>>,
}

impl SegmentReceiverHandle {
    /// Block until the receiver thread exits.
    pub fn join(mut self) {
        if let Some(handle) = self.thread.take() {
            let _ = handle.join();
        }
    }
}

/// Spawn a background thread that receives WAL segments and replays them
/// into the signal ledger.
///
/// The thread exits when `transport.recv_segment()` returns `None` (transport
/// closed / shutdown).
///
/// # Panics
///
/// Panics if the OS fails to spawn the background thread.
pub fn spawn_receiver<T: Transport>(
    transport: Arc<T>,
    ledger: Arc<SignalLedger>,
    replication_state: Arc<ReplicationState>,
) -> SegmentReceiverHandle {
    let thread = std::thread::Builder::new()
        .name("tidaldb-segment-receiver".into())
        .spawn(move || {
            loop {
                let Some(payload) = transport.recv_segment() else {
                    tracing::debug!("segment receiver: transport closed, shutting down");
                    return;
                };

                let shard_id = payload.id.shard_id;
                apply_payload(&payload.bytes, shard_id, &ledger, &replication_state);
            }
        })
        .expect("failed to spawn segment receiver thread");

    SegmentReceiverHandle {
        thread: Some(thread),
    }
}

/// Apply all batches in a WAL segment payload to the signal ledger.
///
/// Idempotent: batches whose last sequence number is at or below the
/// replication state's high-water-mark for the source shard are skipped.
fn apply_payload(
    bytes: &[u8],
    from_shard: ShardId,
    ledger: &SignalLedger,
    state: &ReplicationState,
) {
    let mut offset = 0;
    while offset < bytes.len() {
        let remaining = &bytes[offset..];
        if remaining.len() < HEADER_SIZE {
            break;
        }
        match decode_batch(remaining) {
            Ok((header, events)) => {
                let batch_last_seq = if header.event_count > 0 {
                    header.first_seq + u64::from(header.event_count) - 1
                } else {
                    header.first_seq
                };

                // Idempotency: skip if already applied.
                if let Some(applied) = state.applied_seqno(from_shard)
                    && batch_last_seq <= applied
                {
                    let batch_size = HEADER_SIZE + header.payload_len as usize;
                    offset += batch_size;
                    continue;
                }

                // Apply each event to the ledger.
                for event in &events {
                    let signal_type_id = SignalTypeId::new(u16::from(event.signal_type));
                    let entity_id = EntityId::new(event.entity_id);
                    let weight = f64::from(event.weight);
                    let timestamp = Timestamp::from_nanos(event.timestamp_nanos);
                    ledger.apply_wal_event(signal_type_id, entity_id, weight, timestamp);
                }

                // Advance replication high-water-mark.
                state.advance(from_shard, batch_last_seq);

                let batch_size = HEADER_SIZE + header.payload_len as usize;
                offset += batch_size;
            }
            Err(e) => {
                tracing::warn!(
                    error = %e,
                    "receiver: corrupt batch, skipping remainder of payload"
                );
                break;
            }
        }
    }
}

#[cfg(test)]
#[allow(clippy::unwrap_used)]
mod tests {
    use std::time::Duration;

    use super::*;
    use crate::replication::segment_id::WalSegmentId;
    use crate::replication::shard::RegionId;
    use crate::schema::{DecaySpec, SchemaBuilder, Window};
    use crate::signals::NoopWalWriter;
    use crate::wal::format::batch::{EventRecord, encode_batch};

    fn make_schema() -> crate::schema::Schema {
        let mut builder = SchemaBuilder::new();
        let _ = builder
            .signal(
                "view",
                crate::schema::EntityKind::Item,
                DecaySpec::Exponential {
                    half_life: Duration::from_secs(7 * 24 * 3600),
                },
            )
            .windows(&[Window::AllTime])
            .velocity(false)
            .add();
        builder.build().unwrap()
    }

    fn make_event(entity_id: u64, signal_type: u8, ts_ns: u64) -> EventRecord {
        EventRecord {
            entity_id,
            signal_type,
            weight: 1.0,
            timestamp_nanos: ts_ns,
        }
    }

    #[test]
    fn apply_payload_updates_ledger() {
        let schema = make_schema();
        let ledger = Arc::new(SignalLedger::new(schema, Box::new(NoopWalWriter)));
        let state = Arc::new(ReplicationState::new(&[ShardId::SINGLE]));

        // Resolve the signal type for "view" to get the correct type id.
        let type_id = ledger.resolve_signal_type("view").unwrap();

        let events = vec![make_event(42, type_id.as_u16() as u8, 1_000_000_000)];
        let bytes = encode_batch(&events, 1, 1).unwrap();

        apply_payload(&bytes, ShardId::SINGLE, &ledger, &state);

        // Verify the ledger was updated.
        assert!(ledger.entries().contains_key(&(EntityId::new(42), type_id)));
        assert_eq!(state.applied_seqno(ShardId::SINGLE), Some(1));
    }

    #[test]
    fn apply_payload_idempotent() {
        let schema = make_schema();
        let ledger = Arc::new(SignalLedger::new(schema, Box::new(NoopWalWriter)));
        let state = Arc::new(ReplicationState::new(&[ShardId::SINGLE]));

        let type_id = ledger.resolve_signal_type("view").unwrap();
        let events = vec![make_event(42, type_id.as_u16() as u8, 1_000_000_000)];
        let bytes = encode_batch(&events, 1, 1).unwrap();

        // Apply once.
        apply_payload(&bytes, ShardId::SINGLE, &ledger, &state);
        // Apply again -- should be idempotent.
        apply_payload(&bytes, ShardId::SINGLE, &ledger, &state);

        assert_eq!(state.applied_seqno(ShardId::SINGLE), Some(1));
    }

    #[test]
    fn apply_payload_multiple_batches() {
        let schema = make_schema();
        let ledger = Arc::new(SignalLedger::new(schema, Box::new(NoopWalWriter)));
        let state = Arc::new(ReplicationState::new(&[ShardId::SINGLE]));

        let type_id = ledger.resolve_signal_type("view").unwrap();
        let e1 = vec![make_event(1, type_id.as_u16() as u8, 100)];
        let e2 = vec![make_event(2, type_id.as_u16() as u8, 200)];

        let mut bytes = encode_batch(&e1, 1, 100).unwrap();
        bytes.extend(encode_batch(&e2, 2, 200).unwrap());

        apply_payload(&bytes, ShardId::SINGLE, &ledger, &state);

        assert!(ledger.entries().contains_key(&(EntityId::new(1), type_id)));
        assert!(ledger.entries().contains_key(&(EntityId::new(2), type_id)));
        assert_eq!(state.applied_seqno(ShardId::SINGLE), Some(2));
    }

    /// A minimal transport that returns one payload then signals shutdown.
    struct OneShot {
        rx: crossbeam::channel::Receiver<crate::replication::WalSegmentPayload>,
    }

    impl Transport for OneShot {
        fn send_segment(
            &self,
            _to: ShardId,
            _payload: crate::replication::WalSegmentPayload,
        ) -> Result<(), crate::replication::TransportError> {
            Ok(())
        }
        fn recv_segment(&self) -> Option<crate::replication::WalSegmentPayload> {
            self.rx.recv().ok()
        }
        fn local_shard(&self) -> ShardId {
            ShardId(1)
        }
    }

    #[test]
    fn receiver_thread_exits_on_transport_close() {
        let (tx, rx) = crossbeam::channel::bounded(4);
        let transport = Arc::new(OneShot { rx });

        let schema = make_schema();
        let ledger = Arc::new(SignalLedger::new(schema, Box::new(NoopWalWriter)));
        let state = Arc::new(ReplicationState::new(&[ShardId(0)]));

        let handle = spawn_receiver(
            Arc::clone(&transport),
            Arc::clone(&ledger),
            Arc::clone(&state),
        );

        // Send one segment.
        let type_id = ledger.resolve_signal_type("view").unwrap();
        let events = vec![make_event(99, type_id.as_u16() as u8, 100)];
        let payload_bytes = encode_batch(&events, 1, 1).unwrap();
        tx.send(crate::replication::WalSegmentPayload {
            id: WalSegmentId::new(RegionId::SINGLE, ShardId(0), 1),
            bytes: payload_bytes,
            event_count: 1,
        })
        .unwrap();

        // Give the receiver a moment to process.
        std::thread::sleep(Duration::from_millis(50));

        // Drop sender -- receiver's recv will return None.
        drop(tx);

        // The receiver should exit gracefully.
        handle.join();

        // Verify the segment was applied.
        assert!(ledger.entries().contains_key(&(EntityId::new(99), type_id)));
        assert_eq!(state.applied_seqno(ShardId(0)), Some(1));
    }
}
