//! Session replication bridge.
//!
//! Replicates session journal entries to follower nodes on a dedicated
//! crossbeam channel pair, independent of the WAL segment transport.
//!
//! # Architecture
//!
//! Session events live on a separate transport from WAL segments because:
//! 1. Session writes are latency-sensitive (interactive user sessions).
//! 2. WAL segments are batch-shipped on seal -- different cadence.
//! 3. Independent flow control: a slow WAL receiver must not block session
//!    replication and vice versa.
//!
//! The bridge provides two layers of deduplication on the receiver side:
//! - **Layer 1 (`SeqNo` HWM):** `SessionSeqNoTracker` rejects events with
//!   `seqno <= hwm` for that session. This is the primary defense and is
//!   O(1) per event.
//! - **Layer 2 (Idempotency key):** `IdempotencyStore` (bounded LRU) catches
//!   duplicates that arrive with a fresh seqno due to sender retries after
//!   ambiguous failures.

use std::collections::HashMap;
use std::sync::Arc;

use blake3::Hasher;
use crossbeam::channel::{Receiver, Sender, TrySendError, bounded};
use dashmap::DashMap;

use crate::replication::idempotency::{IdempotencyKey, IdempotencyStore};
use crate::replication::shard::ShardId;
use crate::session::state::SessionSeqNoTracker;
use crate::wal::format::session::{
    SessionSeqNo, SessionWalEvent, decode_session_events, encode_session_event,
};

/// Envelope for session events shipped between shards.
///
/// Uses the existing binary codec (`encode_session_event`) for the payload,
/// not JSON or bincode -- avoids adding new serialization deps.
#[derive(Debug, Clone)]
pub struct SessionPayload {
    /// Source shard that produced these events.
    pub source_shard: ShardId,
    /// Binary-encoded session events (using `encode_session_event`).
    pub bytes: Vec<u8>,
    /// BLAKE3 checksum of `bytes` for integrity verification.
    pub checksum: [u8; 32],
    /// Number of events in this batch.
    pub event_count: u32,
}

impl SessionPayload {
    /// Build a payload from a list of session events.
    ///
    /// Encodes each event using the existing `encode_session_event` codec,
    /// concatenates the results, and computes a BLAKE3 checksum over the
    /// entire byte buffer. The checksum is verified by the receiver before
    /// any events are decoded.
    #[must_use]
    pub fn build(source_shard: ShardId, events: &[SessionWalEvent]) -> Self {
        let mut bytes = Vec::new();
        for event in events {
            bytes.extend(encode_session_event(event));
        }
        let checksum = *Hasher::new().update(&bytes).finalize().as_bytes();
        #[allow(clippy::cast_possible_truncation)]
        Self {
            source_shard,
            bytes,
            checksum,
            event_count: events.len() as u32,
        }
    }

    /// Decode and verify the events in this payload.
    ///
    /// # Errors
    ///
    /// Returns `SessionBridgeError::ChecksumMismatch` if the BLAKE3 checksum
    /// of `bytes` does not match the stored `checksum`.
    pub fn decode_and_verify(&self) -> Result<Vec<SessionWalEvent>, SessionBridgeError> {
        let expected = *Hasher::new().update(&self.bytes).finalize().as_bytes();
        if expected != self.checksum {
            return Err(SessionBridgeError::ChecksumMismatch);
        }
        Ok(decode_session_events(&self.bytes))
    }
}

/// Errors from the session replication bridge.
#[derive(Debug, thiserror::Error)]
pub enum SessionBridgeError {
    /// The BLAKE3 checksum of the payload bytes does not match the
    /// checksum in the envelope header. The payload is corrupted and
    /// must be discarded.
    #[error("session batch checksum mismatch")]
    ChecksumMismatch,
    /// The target shard ID was not registered with the transport factory.
    #[error("unknown peer shard: {0}")]
    UnknownPeer(ShardId),
    /// The session transport channel is full or disconnected.
    #[error("session transport channel closed")]
    Closed,
}

/// In-process session transport factory.
///
/// Creates one sender+receiver pair per shard for session payloads,
/// separate from the WAL segment transport channels.
pub struct InProcessSessionTransportFactory {
    shards: Vec<ShardId>,
    capacity: usize,
}

impl InProcessSessionTransportFactory {
    /// Create a factory for the given shard IDs with default capacity (64).
    #[must_use]
    pub fn new(shards: &[ShardId]) -> Self {
        Self {
            shards: shards.to_vec(),
            capacity: 64,
        }
    }

    /// Create a factory with explicit channel capacity.
    #[must_use]
    pub fn with_capacity(shards: &[ShardId], capacity: usize) -> Self {
        Self {
            shards: shards.to_vec(),
            capacity,
        }
    }

    /// Build one [`SessionShardTransport`] per shard.
    ///
    /// Each transport owns all senders (one per shard) and its own unique
    /// receiver. This mirrors the pattern in
    /// [`super::in_process::InProcessTransportFactory`].
    ///
    /// # Panics
    ///
    /// Panics if a receiver is missing for a shard -- unreachable because
    /// the receiver map is populated from the same shard list.
    #[must_use]
    pub fn build(self) -> HashMap<ShardId, SessionShardTransport> {
        let mut senders: HashMap<ShardId, Sender<SessionPayload>> = HashMap::new();
        let mut receivers: HashMap<ShardId, Receiver<SessionPayload>> = HashMap::new();

        for &shard_id in &self.shards {
            let (tx, rx) = bounded(self.capacity);
            senders.insert(shard_id, tx);
            receivers.insert(shard_id, rx);
        }

        let mut result = HashMap::new();
        for &local_shard in &self.shards {
            let receiver = receivers
                .remove(&local_shard)
                .expect("receiver missing for shard that was just created");
            result.insert(
                local_shard,
                SessionShardTransport {
                    local_shard,
                    senders: senders.clone(),
                    receiver,
                },
            );
        }
        result
    }
}

/// One end of an in-process session transport.
///
/// Holds all senders (for shipping to any peer) and its own receiver
/// (for receiving from any peer). Non-blocking send, blocking receive.
pub struct SessionShardTransport {
    local_shard: ShardId,
    senders: HashMap<ShardId, Sender<SessionPayload>>,
    receiver: Receiver<SessionPayload>,
}

impl SessionShardTransport {
    /// Send a session payload to the `to` shard (non-blocking, best-effort).
    ///
    /// # Errors
    ///
    /// - `SessionBridgeError::UnknownPeer` if `to` is not a registered shard.
    /// - `SessionBridgeError::Closed` if the channel is full or disconnected.
    pub fn send(&self, to: ShardId, payload: SessionPayload) -> Result<(), SessionBridgeError> {
        let tx = self
            .senders
            .get(&to)
            .ok_or(SessionBridgeError::UnknownPeer(to))?;
        tx.try_send(payload).map_err(|e| match e {
            TrySendError::Full(_) | TrySendError::Disconnected(_) => SessionBridgeError::Closed,
        })
    }

    /// Block until a session payload arrives, or return `None` on shutdown.
    ///
    /// Returns `None` when all senders have been dropped (transport closed).
    /// The receiver thread should exit when this returns `None`.
    #[must_use]
    pub fn recv(&self) -> Option<SessionPayload> {
        self.receiver.recv().ok()
    }

    /// The shard identity of this transport endpoint.
    #[must_use]
    pub const fn local_shard(&self) -> ShardId {
        self.local_shard
    }
}

/// Session replication bridge.
///
/// Manages shipping un-shipped session journal entries to peer shards and
/// applying incoming session payloads with idempotency guarantees.
///
/// # Ship side
///
/// `ship()` tracks a per-`(session_id, target_shard)` high-water mark and
/// only ships events with `seqno > hwm`. After a successful send the HWM
/// is advanced, so the same events are never re-shipped unless the bridge
/// is reconstructed (e.g., after crash recovery).
///
/// # Receive side
///
/// `recv_and_apply()` validates the checksum, runs two layers of dedup
/// (`SeqNo` HWM + idempotency key), then applies each surviving event via
/// a caller-provided callback.
pub struct SessionReplicationBridge {
    local_shard: ShardId,
    transport: SessionShardTransport,
    idempotency_store: Arc<IdempotencyStore>,
    seqno_tracker: Arc<SessionSeqNoTracker>,
    /// `(session_id_raw, target_shard)` -> highest seqno shipped.
    ship_hwm: DashMap<(u64, ShardId), SessionSeqNo>,
}

impl SessionReplicationBridge {
    /// Create a new bridge with the given transport and dedup state.
    pub fn new(
        transport: SessionShardTransport,
        idempotency_store: Arc<IdempotencyStore>,
        seqno_tracker: Arc<SessionSeqNoTracker>,
    ) -> Self {
        let local_shard = transport.local_shard();
        Self {
            local_shard,
            transport,
            idempotency_store,
            seqno_tracker,
            ship_hwm: DashMap::new(),
        }
    }

    /// Ship a batch of session events to `target`.
    ///
    /// Only events with `session_seqno > current ship HWM` are included.
    /// Updates ship HWM on success. Does nothing if the filtered batch is
    /// empty (returns `Ok(0)`).
    ///
    /// Legacy events (with `session_seqno == None`) are always shipped
    /// because they predate the seqno mechanism.
    ///
    /// # Errors
    ///
    /// Returns `SessionBridgeError::UnknownPeer` or `SessionBridgeError::Closed`
    /// if the transport send fails.
    pub fn ship(
        &self,
        target: ShardId,
        session_id_raw: u64,
        events: &[SessionWalEvent],
    ) -> Result<usize, SessionBridgeError> {
        let hwm_key = (session_id_raw, target);
        let current_hwm = self
            .ship_hwm
            .get(&hwm_key)
            .map_or(SessionSeqNo::ZERO, |v| *v);

        // Filter: only events with seqno > current HWM (or legacy with no seqno).
        let to_ship: Vec<SessionWalEvent> = events
            .iter()
            .filter(|e| {
                let seqno = extract_seqno(e);
                seqno.is_none_or(|s| s > current_hwm)
            })
            .cloned()
            .collect();

        if to_ship.is_empty() {
            return Ok(0);
        }

        let highest = to_ship
            .iter()
            .filter_map(extract_seqno)
            .max()
            .unwrap_or(current_hwm);

        let count = to_ship.len();
        let payload = SessionPayload::build(self.local_shard, &to_ship);
        self.transport.send(target, payload)?;

        self.ship_hwm.insert(hwm_key, highest);
        Ok(count)
    }

    /// Receive and apply one incoming session payload.
    ///
    /// Validates checksum, applies idempotency checks (seqno HWM + key),
    /// then applies each new event via the provided callback.
    ///
    /// Returns the number of events actually applied (after deduplication).
    /// Returns `None` if the transport is shut down.
    pub fn recv_and_apply(&self, mut apply_fn: impl FnMut(&SessionWalEvent)) -> Option<usize> {
        let payload = self.transport.recv()?;

        let events = match payload.decode_and_verify() {
            Ok(e) => e,
            Err(e) => {
                // Log and skip corrupted payloads.
                tracing::warn!(
                    "session_bridge: corrupted payload from {}: {e}",
                    payload.source_shard
                );
                return Some(0);
            }
        };

        let mut applied = 0;
        for event in &events {
            // Layer 1: SeqNo HWM check.
            if let SessionWalEvent::Signal {
                session_id,
                session_seqno: Some(seqno),
                ..
            } = event
                && !self.seqno_tracker.should_apply(*session_id, *seqno)
            {
                continue;
            }

            // Layer 2: Idempotency key check.
            // Use the stored key directly — the sender already derived it.
            // Note: separate from Layer 1 above (different pattern bindings, different dedup layer).
            #[allow(clippy::collapsible_if)]
            if let SessionWalEvent::Signal {
                idempotency_key: Some(key_int),
                ..
            } = event
                && !self
                    .idempotency_store
                    .check_and_record(IdempotencyKey(*key_int))
            {
                continue;
            }

            apply_fn(event);
            applied += 1;
        }
        Some(applied)
    }

    /// The local shard identity of this bridge.
    #[must_use]
    pub const fn local_shard(&self) -> ShardId {
        self.local_shard
    }

    /// Access the seqno tracker (for test assertions).
    #[must_use]
    pub fn seqno_tracker(&self) -> &SessionSeqNoTracker {
        &self.seqno_tracker
    }
}

/// Extract the session seqno from a `SessionWalEvent`, if present.
const fn extract_seqno(event: &SessionWalEvent) -> Option<SessionSeqNo> {
    match event {
        SessionWalEvent::Signal { session_seqno, .. } => *session_seqno,
        SessionWalEvent::Start { .. } | SessionWalEvent::Close { .. } => None,
    }
}

#[cfg(test)]
#[allow(clippy::unwrap_used)]
mod tests {
    use super::*;

    fn two_shards() -> (SessionReplicationBridge, SessionReplicationBridge) {
        let shards = [ShardId(0), ShardId(1)];
        let mut transports = InProcessSessionTransportFactory::new(&shards).build();
        let t0 = transports.remove(&ShardId(0)).unwrap();
        let t1 = transports.remove(&ShardId(1)).unwrap();
        let store0 = Arc::new(IdempotencyStore::new(1000));
        let store1 = Arc::new(IdempotencyStore::new(1000));
        let tracker0 = Arc::new(SessionSeqNoTracker::new());
        let tracker1 = Arc::new(SessionSeqNoTracker::new());
        let b0 = SessionReplicationBridge::new(t0, store0, tracker0);
        let b1 = SessionReplicationBridge::new(t1, store1, tracker1);
        (b0, b1)
    }

    fn signal_event(session_id: u64, seqno: u64) -> SessionWalEvent {
        SessionWalEvent::Signal {
            session_id,
            entity_id: 42,
            weight: 1.0,
            ts_ns: 1_000_000,
            signal_name: "view".to_string(),
            annotation: None,
            session_seqno: Some(SessionSeqNo(seqno)),
            idempotency_key: Some(seqno as u128 * 1000 + session_id as u128),
        }
    }

    #[test]
    fn ship_and_receive_10_events() {
        let (b0, b1) = two_shards();
        let events: Vec<_> = (1..=10).map(|i| signal_event(1, i)).collect();

        let shipped = b0.ship(ShardId(1), 1, &events).unwrap();
        assert_eq!(shipped, 10);

        let mut received = Vec::new();
        let applied = b1.recv_and_apply(|e| received.push(e.clone())).unwrap();
        assert_eq!(applied, 10);
        assert_eq!(received.len(), 10);
    }

    #[test]
    fn duplicate_payload_produces_zero_shipped() {
        let (b0, _b1) = two_shards();
        let events = vec![signal_event(2, 1)];

        b0.ship(ShardId(1), 2, &events).unwrap();
        let applied_first = _b1.recv_and_apply(|_| {}).unwrap();
        assert_eq!(applied_first, 1);

        // Ship the same events again.
        // Since ship_hwm is advanced, b0 skips them.
        let shipped_again = b0.ship(ShardId(1), 2, &events).unwrap();
        assert_eq!(shipped_again, 0, "ship HWM prevents re-shipping");
    }

    #[test]
    fn checksum_mismatch_returns_zero() {
        let shards = [ShardId(0), ShardId(1)];
        let mut transports = InProcessSessionTransportFactory::new(&shards).build();
        let t0 = transports.remove(&ShardId(0)).unwrap();
        let t1 = transports.remove(&ShardId(1)).unwrap();
        let store1 = Arc::new(IdempotencyStore::new(100));
        let tracker1 = Arc::new(SessionSeqNoTracker::new());

        // Send a corrupted payload directly to t1's channel.
        let mut payload = SessionPayload::build(ShardId(0), &[signal_event(3, 1)]);
        payload.checksum[0] ^= 0xFF; // corrupt the checksum
        t0.send(ShardId(1), payload).unwrap();

        let b1 = SessionReplicationBridge::new(t1, store1, tracker1);
        let applied = b1.recv_and_apply(|_| {}).unwrap();
        assert_eq!(applied, 0, "corrupted payload skipped");
    }

    #[test]
    fn seqno_hwm_prevents_duplicate_events_on_receiver() {
        let (b0, b1) = two_shards();

        // Ship seqno 1, 2, 3.
        let events = vec![signal_event(4, 1), signal_event(4, 2), signal_event(4, 3)];
        b0.ship(ShardId(1), 4, &events).unwrap();
        let applied = b1.recv_and_apply(|_| {}).unwrap();
        assert_eq!(applied, 3);

        // Verify HWM advanced to seqno 3 on receiver.
        assert_eq!(b1.seqno_tracker().hwm(4), SessionSeqNo(3));
    }

    #[test]
    fn session_payload_empty_batch_is_noop() {
        let (b0, _b1) = two_shards();
        let shipped = b0.ship(ShardId(1), 5, &[]).unwrap();
        assert_eq!(shipped, 0);
    }

    #[test]
    fn ship_incremental_only_sends_new_events() {
        let (b0, b1) = two_shards();

        // Ship first batch: seqno 1-3.
        let batch1: Vec<_> = (1..=3).map(|i| signal_event(10, i)).collect();
        let shipped1 = b0.ship(ShardId(1), 10, &batch1).unwrap();
        assert_eq!(shipped1, 3);
        let applied1 = b1.recv_and_apply(|_| {}).unwrap();
        assert_eq!(applied1, 3);

        // Ship second batch: seqno 1-5. Only 4 and 5 should be shipped.
        let batch2: Vec<_> = (1..=5).map(|i| signal_event(10, i)).collect();
        let shipped2 = b0.ship(ShardId(1), 10, &batch2).unwrap();
        assert_eq!(shipped2, 2, "only seqno 4 and 5 should be shipped");

        let mut received = Vec::new();
        let applied2 = b1.recv_and_apply(|e| received.push(e.clone())).unwrap();
        assert_eq!(applied2, 2);
    }

    #[test]
    fn payload_build_and_decode_roundtrip() {
        let events: Vec<_> = (1..=5).map(|i| signal_event(1, i)).collect();
        let payload = SessionPayload::build(ShardId(0), &events);

        assert_eq!(payload.event_count, 5);
        assert_eq!(payload.source_shard, ShardId(0));

        let decoded = payload.decode_and_verify().unwrap();
        assert_eq!(decoded.len(), 5);
        for (orig, dec) in events.iter().zip(decoded.iter()) {
            assert_eq!(orig, dec);
        }
    }

    #[test]
    fn payload_decode_corrupt_bytes_fails() {
        let mut payload = SessionPayload::build(ShardId(0), &[signal_event(1, 1)]);
        // Flip a bit in the payload bytes.
        if !payload.bytes.is_empty() {
            payload.bytes[0] ^= 0x01;
        }
        let result = payload.decode_and_verify();
        assert!(result.is_err());
        assert!(matches!(
            result.unwrap_err(),
            SessionBridgeError::ChecksumMismatch
        ));
    }

    #[test]
    fn send_to_unknown_peer_fails() {
        let (b0, _b1) = two_shards();
        let events = vec![signal_event(1, 1)];
        let result = b0.ship(ShardId(99), 1, &events);
        assert!(result.is_err());
        assert!(matches!(
            result.unwrap_err(),
            SessionBridgeError::UnknownPeer(ShardId(99))
        ));
    }

    #[test]
    fn factory_with_capacity_builds_correctly() {
        let shards = [ShardId(0), ShardId(1)];
        let transports = InProcessSessionTransportFactory::with_capacity(&shards, 8).build();
        assert_eq!(transports.len(), 2);
        assert!(transports.contains_key(&ShardId(0)));
        assert!(transports.contains_key(&ShardId(1)));
    }

    #[test]
    fn start_and_close_events_always_applied() {
        let (b0, b1) = two_shards();

        let events = vec![
            SessionWalEvent::Start {
                session_id: 20,
                user_id: 100,
                started_at_ns: 1_000,
                agent_id: "agent".to_string(),
                policy_name: "default".to_string(),
            },
            signal_event(20, 1),
            SessionWalEvent::Close { session_id: 20 },
        ];

        // Manually build and send (Start/Close have no seqno, always shipped).
        let payload = SessionPayload::build(b0.local_shard, &events);
        b0.transport.send(ShardId(1), payload).unwrap();

        let mut received = Vec::new();
        let applied = b1.recv_and_apply(|e| received.push(e.clone())).unwrap();
        assert_eq!(applied, 3, "Start + Signal + Close all applied");
    }
}
