//! 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`].
//!
//! # Error handling
//!
//! If a batch fails BLAKE3 verification or structural decode, `apply_payload`
//! returns [`WalError::Corruption`]. The receiver thread propagates this
//! error immediately (it does **not** skip the payload) so that
//! [`SegmentReceiverHandle::join`] can surface it to the operator.

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::error::WalError;
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 and retrieve any
/// corruption error that caused it to stop.
pub struct SegmentReceiverHandle {
    thread: Option<JoinHandle<Result<(), WalError>>>,
}

impl SegmentReceiverHandle {
    /// Block until the receiver thread exits.
    ///
    /// # Errors
    ///
    /// Returns `Err(WalError::Corruption { .. })` if a corrupt batch was
    /// received. Returns `Err(WalError::Corruption { .. })` if the thread
    /// panicked unexpectedly.
    pub fn join(mut self) -> Result<(), WalError> {
        self.thread.take().map_or(Ok(()), |handle| {
            handle.join().unwrap_or_else(|_| {
                Err(WalError::Corruption {
                    message: "segment receiver thread panicked".into(),
                })
            })
        })
    }
}

/// 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) **or** when a corrupt batch is detected.
///
/// # Panics
///
/// Panics if the OS fails to spawn the background thread.
pub fn spawn_receiver<T: Transport + ?Sized>(
    transport: Arc<T>,
    ledger: Arc<SignalLedger>,
    replication_state: Arc<ReplicationState>,
) -> SegmentReceiverHandle {
    let thread = std::thread::Builder::new()
        .name("tidaldb-segment-receiver".into())
        .spawn(move || -> Result<(), WalError> {
            loop {
                let Some(payload) = transport.recv_segment() else {
                    tracing::debug!("segment receiver: transport closed, shutting down");
                    return Ok(());
                };

                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.
///
/// # Errors
///
/// Returns `WalError::Corruption` on any BLAKE3 or structural decode failure.
/// The offset of the first corrupt batch is included in the error message.
fn apply_payload(
    bytes: &[u8],
    from_shard: ShardId,
    ledger: &SignalLedger,
    state: &ReplicationState,
) -> Result<(), WalError> {
    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) => {
                return Err(WalError::Corruption {
                    message: format!("corrupt batch at payload offset {offset}: {e}"),
                });
            }
        }
    }
    Ok(())
}

#[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).unwrap();

        // 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).unwrap();
        // Apply again -- should be idempotent.
        apply_payload(&bytes, ShardId::SINGLE, &ledger, &state).unwrap();

        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).unwrap();

        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));
    }

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

        // Build a valid batch then corrupt the checksum region.
        let type_id = ledger.resolve_signal_type("view").unwrap();
        let events = vec![make_event(7, type_id.as_u16() as u8, 500)];
        let mut bytes = encode_batch(&events, 1, 1).unwrap();
        // Flip bytes in the checksum region (bytes 32-63) to break BLAKE3.
        for b in bytes[32..64].iter_mut() {
            *b = b.wrapping_add(1);
        }

        let result = apply_payload(&bytes, ShardId::SINGLE, &ledger, &state);
        assert!(
            matches!(result, Err(WalError::Corruption { .. })),
            "expected Corruption, got {result:?}"
        );
        // Ledger must NOT have been updated.
        assert!(!ledger.entries().contains_key(&(EntityId::new(7), type_id)));
    }

    /// 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().unwrap();

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

    #[test]
    fn receiver_thread_exits_on_corrupt_payload() {
        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),
        );

        // Build a valid batch then corrupt it.
        let type_id = ledger.resolve_signal_type("view").unwrap();
        let events = vec![make_event(55, type_id.as_u16() as u8, 100)];
        let mut corrupt_bytes = encode_batch(&events, 1, 1).unwrap();
        for b in corrupt_bytes[32..64].iter_mut() {
            *b = b.wrapping_add(1);
        }

        tx.send(crate::replication::WalSegmentPayload {
            id: WalSegmentId::new(RegionId::SINGLE, ShardId(0), 1),
            bytes: corrupt_bytes,
            event_count: 1,
        })
        .unwrap();

        // Give thread time to process.
        std::thread::sleep(Duration::from_millis(50));
        drop(tx);

        let result = handle.join();
        assert!(
            matches!(result, Err(WalError::Corruption { .. })),
            "expected Corruption from join, got {result:?}"
        );
        // Entity 55 must NOT have been applied.
        assert!(!ledger.entries().contains_key(&(EntityId::new(55), type_id)));
    }
}
