//! The SFU (Selective Forwarding Unit) media loop.
//!
//! One thread owns every peer's str0m `Rtc` plus the single shared UDP socket
//! (peers are demultiplexed via `rtc.accepts()`). Media is forwarded between
//! peers of the same room WITHOUT decoding. Commands arrive from the control
//! socket over an mpsc channel; notices flow back the other way.
//!
//! The negotiation model follows str0m's canonical SFU example (`chat.rs`):
//! the initial browser offer/answer goes through the control plane (PHP
//! signaling); every renegotiation after that (new tracks as peers join) is
//! exchanged over the WebRTC DATA CHANNEL directly — PHP never relays it.

use std::collections::VecDeque;
use std::io::ErrorKind;
use std::net::{SocketAddr, UdpSocket};
use std::ops::Deref;
use std::path::Path;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::mpsc::{Receiver, Sender, TryRecvError};
use std::sync::{Arc, Weak};
use std::time::{Duration, Instant};

use log::{debug, info, warn};
use str0m::change::{SdpAnswer, SdpOffer, SdpPendingOffer};
use str0m::channel::{ChannelData, ChannelId};
use str0m::media::{Direction, KeyframeRequest, KeyframeRequestKind, MediaData, MediaKind, Mid, Rid};
use str0m::net::{Protocol, Receive};
use str0m::{Candidate, Event, IceConnectionState, Input, Output, Rtc};

use crate::proto::{Command, Notice, PeerStats, Reply, Request, RoomStats};
use crate::recording::OggOpusWriter;

/// A control request paired with where to send its reply.
pub struct CtrlMsg {
    pub request: Request,
    pub reply: Sender<Reply>,
}

pub struct Sfu {
    socket: UdpSocket,
    /// The candidate address advertised to browsers (public ip + UDP port).
    advertise: SocketAddr,
    commands: Receiver<CtrlMsg>,
    notices: Sender<Notice>,
    clients: Vec<Client>,
}

impl Sfu {
    pub fn new(
        socket: UdpSocket,
        advertise: SocketAddr,
        commands: Receiver<CtrlMsg>,
        notices: Sender<Notice>,
    ) -> Self {
        Sfu {
            socket,
            advertise,
            commands,
            notices,
            clients: Vec::new(),
        }
    }

    /// Run until a `shutdown` command arrives.
    pub fn run(mut self) {
        let mut to_propagate: VecDeque<Propagated> = VecDeque::new();
        let mut buf = vec![0; 2000];

        loop {
            self.reap_dead_clients();

            match self.drain_commands() {
                ControlFlow::Continue => {}
                ControlFlow::Shutdown => {
                    info!("shutdown command received, stopping SFU loop");
                    return;
                }
            }

            // Poll clients until they all return a timeout.
            let mut timeout = Instant::now() + Duration::from_millis(100);
            for client in self.clients.iter_mut() {
                let t = poll_until_timeout(client, &mut to_propagate, &self.socket);
                timeout = timeout.min(t);
            }

            // If we have an item to propagate, do that before sleeping.
            if let Some(p) = to_propagate.pop_front() {
                self.propagate(p);
                continue;
            }

            let duration = (timeout - Instant::now()).max(Duration::from_millis(1));
            self.socket
                .set_read_timeout(Some(duration))
                .expect("setting socket read timeout");

            if let Some(input) = read_socket_input(&self.socket, &mut buf) {
                // rtc.accepts() demultiplexes the shared UDP socket to the right peer.
                if let Some(client) = self.clients.iter_mut().find(|c| c.accepts(&input)) {
                    client.handle_input(input);
                } else {
                    // Common: first STUN can arrive before add_peer completes.
                    debug!("no client accepts UDP input: {:?}", input);
                }
            }

            // Drive time forward in all clients.
            let now = Instant::now();
            for client in &mut self.clients {
                client.handle_input(Input::Timeout(now));
            }

            self.notify_connected();
        }
    }

    fn drain_commands(&mut self) -> ControlFlow {
        loop {
            let msg = match self.commands.try_recv() {
                Ok(msg) => msg,
                Err(TryRecvError::Empty) => return ControlFlow::Continue,
                Err(TryRecvError::Disconnected) => {
                    warn!("control channel disconnected, stopping");
                    return ControlFlow::Shutdown;
                }
            };

            let id = msg.request.id;
            let reply = match msg.request.command {
                Command::Ping => Reply::ok(id),
                Command::Shutdown => {
                    let _ = msg.reply.send(Reply::ok(id));
                    return ControlFlow::Shutdown;
                }
                Command::AddPeer {
                    room,
                    peer,
                    offer,
                    record,
                } => self.add_peer(id, room, peer, &offer, record.as_deref()),
                Command::RemovePeer { room, peer } => self.with_client(id, &room, &peer, |c| {
                    c.rtc.disconnect();
                    Reply::ok(id)
                }),
                Command::MutePeer { room, peer, muted } => {
                    self.with_client(id, &room, &peer, |c| {
                        c.muted = muted;
                        Reply::ok(id)
                    })
                }
                Command::RecordStart { room, peer, path } => {
                    self.with_client(id, &room, &peer, |c| match c.start_recording(&path) {
                        Ok(()) => Reply::ok(id),
                        Err(e) => Reply::error(id, format!("recording failed: {e}")),
                    })
                }
                Command::RecordStop { room, peer } => self.with_client(id, &room, &peer, |c| {
                    c.stop_recording();
                    Reply::ok(id)
                }),
                Command::Stats => Reply::stats(id, self.stats()),
            };

            let _ = msg.reply.send(reply);
        }
    }

    fn add_peer(
        &mut self,
        id: u64,
        room: String,
        peer: String,
        offer_sdp: &str,
        record: Option<&str>,
    ) -> Reply {
        let offer = match SdpOffer::from_sdp_string(offer_sdp) {
            Ok(offer) => offer,
            Err(e) => return Reply::error(id, format!("bad offer sdp: {e}")),
        };

        let mut rtc = Rtc::builder().build();

        let candidate = match Candidate::host(self.advertise, "udp") {
            Ok(candidate) => candidate,
            Err(e) => return Reply::error(id, format!("bad host candidate: {e}")),
        };
        rtc.add_local_candidate(candidate);

        let answer = match rtc.sdp_api().accept_offer(offer) {
            Ok(answer) => answer,
            Err(e) => return Reply::error(id, format!("offer rejected: {e}")),
        };

        // A rejoin with the same peer id replaces the old connection.
        if let Some(old) = self
            .clients
            .iter_mut()
            .find(|c| c.room == room && c.peer == peer)
        {
            info!("replacing existing peer {room}/{peer}");
            old.rtc.disconnect();
        }

        let mut client = Client::new(room.clone(), peer.clone(), rtc);
        if let Some(path) = record {
            if let Err(e) = client.start_recording(path) {
                warn!("recording for {room}/{peer} failed to start: {e}");
            }
        }

        // Open the tracks already present in the room towards the new peer.
        for track in self
            .clients
            .iter()
            .filter(|c| c.room == room)
            .flat_map(|c| c.tracks_in.iter())
        {
            client.handle_track_open(Arc::downgrade(&track.id));
        }

        info!("peer {room}/{peer} added ({} in room)", 1 + self
            .clients
            .iter()
            .filter(|c| c.room == room)
            .count());
        self.clients.push(client);

        Reply::answer(id, answer.to_sdp_string())
    }

    fn with_client(
        &mut self,
        id: u64,
        room: &str,
        peer: &str,
        f: impl FnOnce(&mut Client) -> Reply,
    ) -> Reply {
        match self
            .clients
            .iter_mut()
            .find(|c| c.room == room && c.peer == peer)
        {
            Some(client) => f(client),
            None => Reply::error(id, format!("no such peer: {room}/{peer}")),
        }
    }

    fn stats(&self) -> Vec<RoomStats> {
        let mut rooms: Vec<RoomStats> = Vec::new();
        for client in &self.clients {
            let peer = PeerStats {
                peer: client.peer.clone(),
                connected: client.connected_at.is_some(),
                muted: client.muted,
                seconds: client
                    .connected_at
                    .map(|t| t.elapsed().as_secs_f64())
                    .unwrap_or(0.0),
            };
            match rooms.iter_mut().find(|r| r.room == client.room) {
                Some(room) => room.peers.push(peer),
                None => rooms.push(RoomStats {
                    room: client.room.clone(),
                    peers: vec![peer],
                }),
            }
        }
        rooms
    }

    /// Emit `peer_connected` notices for clients that just reached Connected.
    fn notify_connected(&mut self) {
        for client in &mut self.clients {
            if client.connected_at.is_some() && !client.notified_connected {
                client.notified_connected = true;
                let _ = self.notices.send(Notice::PeerConnected {
                    room: client.room.clone(),
                    peer: client.peer.clone(),
                });
            }
        }
    }

    fn reap_dead_clients(&mut self) {
        if self.clients.iter().all(|c| c.rtc.is_alive()) {
            return;
        }

        let clients = std::mem::take(&mut self.clients);
        for mut client in clients {
            if client.rtc.is_alive() {
                self.clients.push(client);
                continue;
            }

            client.stop_recording();
            info!("peer {}/{} left", client.room, client.peer);
            let _ = self.notices.send(Notice::PeerLeft {
                room: client.room.clone(),
                peer: client.peer.clone(),
                seconds: client
                    .connected_at
                    .map(|t| t.elapsed().as_secs_f64())
                    .unwrap_or(0.0),
            });
        }
    }

    /// Send one propagated item to the other clients in the origin's room.
    fn propagate(&mut self, propagated: Propagated) {
        let Some(client_id) = propagated.client_id() else {
            return;
        };

        let Some(origin) = self.clients.iter().find(|c| c.id == client_id) else {
            return; // origin left in the meantime
        };
        let (room, origin_muted) = (origin.room.clone(), origin.muted);

        for client in &mut self.clients {
            if client.id == client_id || client.room != room {
                continue;
            }

            match &propagated {
                Propagated::TrackOpen(_, track_in) => client.handle_track_open(track_in.clone()),
                Propagated::MediaData(_, data) => {
                    if !origin_muted {
                        client.handle_media_data_out(client_id, data)
                    }
                }
                Propagated::KeyframeRequest(_, req, origin_id, mid_in) => {
                    // Only the origin client handles the keyframe request.
                    if *origin_id == client.id {
                        client.handle_keyframe_request(*req, *mid_in)
                    }
                }
                Propagated::Noop | Propagated::Timeout(_) => {}
            }
        }
    }
}

enum ControlFlow {
    Continue,
    Shutdown,
}

/// Poll all the output from the client until it returns a timeout.
fn poll_until_timeout(
    client: &mut Client,
    queue: &mut VecDeque<Propagated>,
    socket: &UdpSocket,
) -> Instant {
    loop {
        if !client.rtc.is_alive() {
            return Instant::now();
        }

        let propagated = client.poll_output(socket);

        if let Propagated::Timeout(t) = propagated {
            return t;
        }

        queue.push_back(propagated)
    }
}

fn read_socket_input<'a>(socket: &UdpSocket, buf: &'a mut Vec<u8>) -> Option<Input<'a>> {
    buf.resize(2000, 0);

    match socket.recv_from(buf) {
        Ok((n, source)) => {
            buf.truncate(n);

            let Ok(contents) = buf.as_slice().try_into() else {
                return None;
            };

            Some(Input::Receive(
                Instant::now(),
                Receive {
                    proto: Protocol::Udp,
                    source,
                    destination: socket.local_addr().unwrap(),
                    contents,
                },
            ))
        }

        Err(e) => match e.kind() {
            ErrorKind::WouldBlock | ErrorKind::TimedOut => None,
            _ => panic!("UdpSocket read failed: {e:?}"),
        },
    }
}

struct Client {
    id: ClientId,
    room: String,
    peer: String,
    rtc: Rtc,
    pending: Option<SdpPendingOffer>,
    cid: Option<ChannelId>,
    tracks_in: Vec<TrackInEntry>,
    tracks_out: Vec<TrackOut>,
    chosen_rid: Option<Rid>,
    muted: bool,
    recorder: Option<OggOpusWriter>,
    connected_at: Option<Instant>,
    notified_connected: bool,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct ClientId(u64);

impl Deref for ClientId {
    type Target = u64;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

struct TrackIn {
    origin: ClientId,
    mid: Mid,
    kind: MediaKind,
}

struct TrackInEntry {
    id: Arc<TrackIn>,
    last_keyframe_request: Option<Instant>,
}

struct TrackOut {
    track_in: Weak<TrackIn>,
    state: TrackOutState,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum TrackOutState {
    ToOpen,
    Negotiating(Mid),
    Open(Mid),
}

impl TrackOut {
    fn mid(&self) -> Option<Mid> {
        match self.state {
            TrackOutState::ToOpen => None,
            TrackOutState::Negotiating(m) | TrackOutState::Open(m) => Some(m),
        }
    }
}

impl Client {
    fn new(room: String, peer: String, rtc: Rtc) -> Client {
        static ID_COUNTER: AtomicU64 = AtomicU64::new(0);
        let next_id = ID_COUNTER.fetch_add(1, Ordering::SeqCst);
        Client {
            id: ClientId(next_id),
            room,
            peer,
            rtc,
            pending: None,
            cid: None,
            tracks_in: vec![],
            tracks_out: vec![],
            chosen_rid: None,
            muted: false,
            recorder: None,
            connected_at: None,
            notified_connected: false,
        }
    }

    fn start_recording(&mut self, path: &str) -> Result<(), std::io::Error> {
        self.recorder = Some(OggOpusWriter::create(Path::new(path), 1)?);
        Ok(())
    }

    fn stop_recording(&mut self) {
        if let Some(mut recorder) = self.recorder.take() {
            let _ = recorder.finish();
        }
    }

    fn accepts(&self, input: &Input) -> bool {
        self.rtc.accepts(input)
    }

    fn handle_input(&mut self, input: Input) {
        if !self.rtc.is_alive() {
            return;
        }

        if let Err(e) = self.rtc.handle_input(input) {
            warn!("client {}/{} disconnected: {:?}", self.room, self.peer, e);
            self.rtc.disconnect();
        }
    }

    fn poll_output(&mut self, socket: &UdpSocket) -> Propagated {
        if !self.rtc.is_alive() {
            return Propagated::Noop;
        }

        // Incoming tracks from room peers cause new entries in tracks_out that
        // need SDP negotiation with this client (over its data channel).
        if self.negotiate_if_needed() {
            return Propagated::Noop;
        }

        match self.rtc.poll_output() {
            Ok(output) => self.handle_output(output, socket),
            Err(e) => {
                warn!(
                    "client {}/{} poll_output failed: {:?}",
                    self.room, self.peer, e
                );
                self.rtc.disconnect();
                Propagated::Noop
            }
        }
    }

    fn handle_output(&mut self, output: Output, socket: &UdpSocket) -> Propagated {
        match output {
            Output::Transmit(transmit) => {
                if let Err(e) = socket.send_to(&transmit.contents, transmit.destination) {
                    warn!("UDP send_to failed: {e:?}");
                }
                Propagated::Noop
            }
            Output::Timeout(t) => Propagated::Timeout(t),
            Output::Event(e) => match e {
                Event::Connected => {
                    self.connected_at = Some(Instant::now());
                    Propagated::Noop
                }
                Event::IceConnectionStateChange(v) => {
                    if v == IceConnectionState::Disconnected {
                        self.rtc.disconnect();
                    }
                    Propagated::Noop
                }
                Event::MediaAdded(e) => self.handle_media_added(e.mid, e.kind),
                Event::MediaData(data) => self.handle_media_data_in(data),
                Event::KeyframeRequest(req) => self.handle_incoming_keyframe_req(req),
                Event::ChannelOpen(cid, _) => {
                    self.cid = Some(cid);
                    Propagated::Noop
                }
                Event::ChannelData(data) => self.handle_channel_data(data),
                _ => Propagated::Noop,
            },
        }
    }

    fn handle_media_added(&mut self, mid: Mid, kind: MediaKind) -> Propagated {
        let track_in = TrackInEntry {
            id: Arc::new(TrackIn {
                origin: self.id,
                mid,
                kind,
            }),
            last_keyframe_request: None,
        };

        let weak = Arc::downgrade(&track_in.id);
        self.tracks_in.push(track_in);

        Propagated::TrackOpen(self.id, weak)
    }

    fn handle_media_data_in(&mut self, data: MediaData) -> Propagated {
        if !data.contiguous {
            self.request_keyframe_throttled(data.mid, data.rid, KeyframeRequestKind::Fir);
        }

        if !self.muted {
            if let Some(recorder) = &mut self.recorder {
                if data.params.spec().codec.is_audio() {
                    if let Err(e) = recorder.write_frame(&data.data, data.time.as_seconds()) {
                        warn!("recording write failed for {}/{}: {e}", self.room, self.peer);
                        self.recorder = None;
                    }
                }
            }
        }

        Propagated::MediaData(self.id, data)
    }

    fn request_keyframe_throttled(
        &mut self,
        mid: Mid,
        rid: Option<Rid>,
        kind: KeyframeRequestKind,
    ) {
        let Some(mut writer) = self.rtc.writer(mid) else {
            return;
        };

        let Some(track_entry) = self.tracks_in.iter_mut().find(|t| t.id.mid == mid) else {
            return;
        };

        if track_entry
            .last_keyframe_request
            .map(|t| t.elapsed() < Duration::from_secs(1))
            .unwrap_or(false)
        {
            return;
        }

        _ = writer.request_keyframe(rid, kind);

        track_entry.last_keyframe_request = Some(Instant::now());
    }

    fn handle_incoming_keyframe_req(&self, mut req: KeyframeRequest) -> Propagated {
        let Some(track_out) = self.tracks_out.iter().find(|t| t.mid() == Some(req.mid)) else {
            return Propagated::Noop;
        };
        let Some(track_in) = track_out.track_in.upgrade() else {
            return Propagated::Noop;
        };

        req.rid = self.chosen_rid;

        Propagated::KeyframeRequest(self.id, req, track_in.origin, track_in.mid)
    }

    fn negotiate_if_needed(&mut self) -> bool {
        if self.cid.is_none() || self.pending.is_some() {
            // No data channel yet, or a negotiation is already in flight.
            return false;
        }

        let mut change = self.rtc.sdp_api();

        for track in &mut self.tracks_out {
            if let TrackOutState::ToOpen = track.state {
                if let Some(track_in) = track.track_in.upgrade() {
                    let stream_id = track_in.origin.to_string();
                    let mid =
                        change.add_media(track_in.kind, Direction::SendOnly, Some(stream_id), None);
                    track.state = TrackOutState::Negotiating(mid);
                }
            }
        }

        if !change.has_changes() {
            return false;
        }

        let Some((offer, pending)) = change.apply() else {
            return false;
        };

        let Some(mut channel) = self.cid.and_then(|id| self.rtc.channel(id)) else {
            return false;
        };

        let json = serde_json::to_string(&offer).expect("offer to serialize");
        if channel.write(false, json.as_bytes()).is_err() {
            return false;
        }

        self.pending = Some(pending);

        true
    }

    fn handle_channel_data(&mut self, d: ChannelData) -> Propagated {
        if let Ok(offer) = serde_json::from_slice::<'_, SdpOffer>(&d.data) {
            self.handle_offer(offer);
        } else if let Ok(answer) = serde_json::from_slice::<'_, SdpAnswer>(&d.data) {
            self.handle_answer(answer);
        }

        Propagated::Noop
    }

    fn handle_offer(&mut self, offer: SdpOffer) {
        let Ok(answer) = self.rtc.sdp_api().accept_offer(offer) else {
            warn!("client {}/{} offer rejected", self.room, self.peer);
            return;
        };

        // Cancel any pending negotiation; it is redone after this offer.
        for track in &mut self.tracks_out {
            if let TrackOutState::Negotiating(_) = track.state {
                track.state = TrackOutState::ToOpen;
            }
        }

        let Some(mut channel) = self.cid.and_then(|id| self.rtc.channel(id)) else {
            return;
        };

        let json = serde_json::to_string(&answer).expect("answer to serialize");
        let _ = channel.write(false, json.as_bytes());
    }

    fn handle_answer(&mut self, answer: SdpAnswer) {
        if let Some(pending) = self.pending.take() {
            if self.rtc.sdp_api().accept_answer(pending, answer).is_err() {
                warn!("client {}/{} answer rejected", self.room, self.peer);
                return;
            }

            for track in &mut self.tracks_out {
                if let TrackOutState::Negotiating(m) = track.state {
                    track.state = TrackOutState::Open(m);
                }
            }
        }
    }

    fn handle_track_open(&mut self, track_in: Weak<TrackIn>) {
        self.tracks_out.push(TrackOut {
            track_in,
            state: TrackOutState::ToOpen,
        });
    }

    fn handle_media_data_out(&mut self, origin: ClientId, data: &MediaData) {
        // Map the incoming media to the outgoing track towards this client.
        let Some(mid) = self
            .tracks_out
            .iter()
            .find(|o| {
                o.track_in
                    .upgrade()
                    .filter(|i| i.origin == origin && i.mid == data.mid)
                    .is_some()
            })
            .and_then(|o| o.mid())
        else {
            return;
        };

        if data.rid.is_some() && data.rid != Some("h".into()) {
            // Simulcast layer selection (video); pass rid=None or the "h" layer.
            return;
        }

        if self.chosen_rid != data.rid {
            self.chosen_rid = data.rid;
        }

        let Some(writer) = self.rtc.writer(mid) else {
            return;
        };

        let Some(pt) = writer.match_params(data.params) else {
            return;
        };

        if let Err(e) = writer.write(pt, data.network_time, data.time, data.data.clone()) {
            warn!("client {}/{} write failed: {:?}", self.room, self.peer, e);
            self.rtc.disconnect();
        }
    }

    fn handle_keyframe_request(&mut self, req: KeyframeRequest, mid_in: Mid) {
        let has_incoming_track = self.tracks_in.iter().any(|i| i.id.mid == mid_in);

        if !has_incoming_track {
            return;
        }

        let Some(mut writer) = self.rtc.writer(mid_in) else {
            return;
        };

        if let Err(e) = writer.request_keyframe(req.rid, req.kind) {
            debug!("request_keyframe failed: {:?}", e);
        }
    }
}

/// Events propagated between clients of a room.
#[allow(clippy::large_enum_variant)]
enum Propagated {
    Noop,
    Timeout(Instant),
    TrackOpen(ClientId, Weak<TrackIn>),
    MediaData(ClientId, MediaData),
    KeyframeRequest(ClientId, KeyframeRequest, ClientId, Mid),
}

impl Propagated {
    fn client_id(&self) -> Option<ClientId> {
        match self {
            Propagated::TrackOpen(c, _)
            | Propagated::MediaData(c, _)
            | Propagated::KeyframeRequest(c, _, _, _) => Some(*c),
            _ => None,
        }
    }
}
