//! SFU-style Selective Forwarding Unit for the Pulse Relay.
//!
//! Routes packets between participants without decoding media.
//! Each participant has a set of published streams and subscribed streams.
//! The SFU forwards packets from publishers to matching subscribers.

use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use tokio::sync::{RwLock, mpsc};
use tracing::{debug, info, warn};

/// Unique identifier for a participant in a session.
pub type ParticipantId = String;
/// Unique identifier for a media/data stream.
pub type StreamId = String;

/// A packet to be forwarded through the SFU.
#[derive(Debug, Clone)]
pub struct ForwardPacket {
    pub source: ParticipantId,
    pub stream_id: StreamId,
    pub sequence: u32,
    pub payload: Vec<u8>,
    pub lossless: bool,
}

/// Represents a connected participant with their publish/subscribe state.
struct Participant {
    id: ParticipantId,
    /// Streams this participant is publishing
    published_streams: HashSet<StreamId>,
    /// Streams this participant wants to receive
    subscribed_streams: HashSet<StreamId>,
    /// Channel to send forwarded packets to this participant
    tx: mpsc::Sender<ForwardPacket>,
}

/// The Selective Forwarding Unit routes packets between participants.
pub struct Sfu {
    participants: Arc<RwLock<HashMap<ParticipantId, Participant>>>,
}

impl Sfu {
    pub fn new() -> Self {
        Self {
            participants: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    /// Register a new participant and return a channel receiver for forwarded packets.
    pub async fn add_participant(&self, id: ParticipantId) -> mpsc::Receiver<ForwardPacket> {
        let (tx, rx) = mpsc::channel(256);

        let participant = Participant {
            id: id.clone(),
            published_streams: HashSet::new(),
            subscribed_streams: HashSet::new(),
            tx,
        };

        self.participants
            .write()
            .await
            .insert(id.clone(), participant);
        info!("[SFU] Participant '{}' joined", id);
        rx
    }

    /// Remove a participant from the SFU.
    pub async fn remove_participant(&self, id: &str) {
        self.participants.write().await.remove(id);
        info!("[SFU] Participant '{}' left", id);
    }

    /// A participant declares they are publishing a stream.
    pub async fn publish(&self, participant_id: &str, stream_id: StreamId) {
        let mut participants = self.participants.write().await;
        if let Some(p) = participants.get_mut(participant_id) {
            p.published_streams.insert(stream_id.clone());
            info!(
                "[SFU] '{}' publishing stream '{}'",
                participant_id, stream_id
            );
        }
    }

    /// A participant subscribes to another participant's stream.
    pub async fn subscribe(&self, participant_id: &str, stream_id: StreamId) {
        let mut participants = self.participants.write().await;
        if let Some(p) = participants.get_mut(participant_id) {
            p.subscribed_streams.insert(stream_id.clone());
            info!(
                "[SFU] '{}' subscribed to stream '{}'",
                participant_id, stream_id
            );
        }
    }

    /// Forward a packet from a publisher to all subscribers of that stream.
    /// This is the hot path — zero media decoding, just route.
    pub async fn forward(&self, packet: ForwardPacket) {
        let participants = self.participants.read().await;

        let mut lossless_sends = Vec::new();

        for (pid, participant) in participants.iter() {
            // Don't send back to the source
            if *pid == packet.source {
                continue;
            }

            // Only forward if this participant is subscribed to the stream
            if participant.subscribed_streams.contains(&packet.stream_id) {
                if packet.lossless {
                    // Queue for slow subscribers
                    lossless_sends.push((participant.tx.clone(), packet.clone()));
                } else {
                    // Drop for slow subscribers (lossy)
                    if let Err(_) = participant.tx.try_send(packet.clone()) {
                        warn!(
                            "[SFU] Dropped packet for '{}' on stream '{}' (backpressure)",
                            pid, packet.stream_id
                        );
                    }
                }
            }
        }

        // Release the read lock before potentially blocking
        drop(participants);

        // Apply backpressure (wait for space) for lossless streams
        for (tx, pkt) in lossless_sends {
            let _ = tx.send(pkt).await;
        }
    }

    /// Returns the number of connected participants.
    pub async fn participant_count(&self) -> usize {
        self.participants.read().await.len()
    }
}

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

    #[tokio::test]
    async fn test_basic_forwarding() {
        let sfu = Sfu::new();

        let mut rx_b = sfu.add_participant("alice".into()).await;
        let mut rx_a = sfu.add_participant("bob".into()).await;

        sfu.publish("alice", "audio".into()).await;
        sfu.subscribe("bob", "audio".into()).await;

        sfu.forward(ForwardPacket {
            source: "alice".into(),
            stream_id: "audio".into(),
            sequence: 1,
            payload: vec![0xDE, 0xAD],
            lossless: false,
        })
        .await;

        let pkt = rx_a.recv().await.expect("bob should receive packet");
        assert_eq!(pkt.sequence, 1);
        assert_eq!(pkt.payload, vec![0xDE, 0xAD]);

        // Alice should NOT receive her own packet
        assert!(rx_b.try_recv().is_err());
    }

    #[tokio::test]
    async fn test_unsubscribed_does_not_receive() {
        let sfu = Sfu::new();

        let mut rx_b = sfu.add_participant("alice".into()).await;
        let mut rx_c = sfu.add_participant("charlie".into()).await;

        sfu.publish("alice", "video".into()).await;
        // charlie does NOT subscribe

        sfu.forward(ForwardPacket {
            source: "alice".into(),
            stream_id: "video".into(),
            sequence: 1,
            payload: vec![1, 2, 3],
            lossless: false,
        })
        .await;

        assert!(
            rx_c.try_recv().is_err(),
            "unsubscribed participant should not receive"
        );
    }

    #[tokio::test]
    async fn test_multi_subscriber() {
        let sfu = Sfu::new();

        let _ = sfu.add_participant("alice".into()).await;
        let mut rx_bob = sfu.add_participant("bob".into()).await;
        let mut rx_charlie = sfu.add_participant("charlie".into()).await;

        sfu.publish("alice", "screen".into()).await;
        sfu.subscribe("bob", "screen".into()).await;
        sfu.subscribe("charlie", "screen".into()).await;

        sfu.forward(ForwardPacket {
            source: "alice".into(),
            stream_id: "screen".into(),
            sequence: 42,
            payload: vec![0xFF],
            lossless: false,
        })
        .await;

        assert_eq!(rx_bob.recv().await.unwrap().sequence, 42);
        assert_eq!(rx_charlie.recv().await.unwrap().sequence, 42);
    }

    #[tokio::test]
    async fn test_backpressure_drop_lossy() {
        let sfu = Sfu::new();
        sfu.add_participant("pub".into()).await;
        // The channel is bounded by 256
        let mut rx_sub = sfu.add_participant("sub".into()).await;
        sfu.publish("pub", "audio".into()).await;
        sfu.subscribe("sub", "audio".into()).await;

        // Fill the channel with 256 packets
        for i in 0..256 {
            sfu.forward(ForwardPacket {
                source: "pub".into(),
                stream_id: "audio".into(),
                sequence: i,
                payload: vec![1],
                lossless: false, // Lossy mode
            })
            .await;
        }

        // Send one more packet when full. This should drop instantly, not hang.
        sfu.forward(ForwardPacket {
            source: "pub".into(),
            stream_id: "audio".into(),
            sequence: 999,
            payload: vec![1],
            lossless: false, // Lossy mode
        })
        .await;

        let first = rx_sub.recv().await.unwrap();
        assert_eq!(first.sequence, 0); // Didn't block, didn't queue the 999th.
    }

    #[tokio::test]
    async fn test_late_joiner_no_replay() {
        let sfu = Sfu::new();
        sfu.add_participant("pub".into()).await;
        sfu.publish("pub", "video".into()).await;

        sfu.forward(ForwardPacket {
            source: "pub".into(),
            stream_id: "video".into(),
            sequence: 1,
            payload: vec![0xA],
            lossless: true,
        })
        .await;

        // Sub joins AFTER packet sent
        let mut rx_sub = sfu.add_participant("sub".into()).await;
        sfu.subscribe("sub", "video".into()).await;

        // Ensure sub got nothing
        assert!(rx_sub.try_recv().is_err());
    }
}
