use super::*;
use crate::transport::PacketBuffer;
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering::Relaxed};
use tokio::sync::mpsc::error::TryRecvError;

#[derive(Debug, Clone)]
pub(crate) struct EndpointServiceDatagramDelivery {
    pub(crate) source_peer: PeerIdentity,
    pub(crate) source_port: u16,
    pub(crate) destination_port: u16,
    pub(crate) payload: PacketBuffer,
    pub(crate) enqueued_at_ms: u64,
}

impl EndpointServiceDatagramDelivery {
    pub(crate) fn new(
        source_peer: PeerIdentity,
        source_port: u16,
        destination_port: u16,
        payload: PacketBuffer,
    ) -> Self {
        Self {
            source_peer,
            source_port,
            destination_port,
            payload,
            enqueued_at_ms: crate::time::now_ms(),
        }
    }
}

#[derive(Debug)]
pub(crate) struct NodeEndpointServiceEvent {
    pub(crate) messages: Vec<EndpointServiceDatagramDelivery>,
}

#[derive(Debug, Clone)]
pub(crate) struct EndpointServiceEventSender {
    tx: tokio::sync::mpsc::UnboundedSender<NodeEndpointServiceEvent>,
    queued_messages: Arc<AtomicUsize>,
    message_capacity: usize,
}

#[derive(Debug)]
pub(crate) struct EndpointServiceEventReceiver {
    rx: tokio::sync::mpsc::UnboundedReceiver<NodeEndpointServiceEvent>,
    queued_messages: Arc<AtomicUsize>,
}

impl EndpointServiceEventSender {
    pub(crate) fn channel(capacity: usize) -> (Self, EndpointServiceEventReceiver) {
        let queued_messages = Arc::new(AtomicUsize::new(0));
        let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
        (
            Self {
                tx,
                queued_messages: Arc::clone(&queued_messages),
                message_capacity: capacity.max(1),
            },
            EndpointServiceEventReceiver {
                rx,
                queued_messages,
            },
        )
    }

    pub(crate) fn send(&self, messages: Vec<EndpointServiceDatagramDelivery>) -> Result<(), ()> {
        if messages.is_empty() {
            return Ok(());
        }
        self.send_messages(messages)
    }

    fn send_messages(&self, mut messages: Vec<EndpointServiceDatagramDelivery>) -> Result<(), ()> {
        let count = messages.len();
        let reserved = self
            .queued_messages
            .fetch_update(Relaxed, Relaxed, |current| {
                current
                    .checked_add(count)
                    .filter(|next| *next <= self.message_capacity)
            })
            .is_ok();
        if !reserved {
            if messages.len() > 1 {
                let right = messages.split_off(messages.len() / 2);
                self.send_messages(messages)?;
                self.send_messages(right)?;
            }
            return Ok(());
        }

        if self.tx.send(NodeEndpointServiceEvent { messages }).is_err() {
            self.queued_messages.fetch_sub(count, Relaxed);
            return Err(());
        }
        Ok(())
    }
}

impl EndpointServiceEventReceiver {
    pub(crate) async fn recv(&mut self) -> Option<NodeEndpointServiceEvent> {
        self.rx.recv().await
    }

    pub(crate) fn blocking_recv(&mut self) -> Option<NodeEndpointServiceEvent> {
        self.rx.blocking_recv()
    }

    pub(crate) fn try_recv(&mut self) -> Result<NodeEndpointServiceEvent, TryRecvError> {
        self.rx.try_recv()
    }

    pub(crate) fn release_messages(&self, count: usize) {
        if count > 0 {
            let previous = self.queued_messages.fetch_sub(count, Relaxed);
            debug_assert!(previous >= count, "service event accounting underflow");
        }
    }
}

impl Drop for EndpointServiceEventReceiver {
    fn drop(&mut self) {
        self.queued_messages.store(0, Relaxed);
    }
}

#[derive(Debug, Default)]
pub(in crate::node) struct EndpointServiceRuntime {
    senders: HashMap<u16, EndpointServiceEventSender>,
}

impl EndpointServiceRuntime {
    pub(in crate::node) fn register(
        &mut self,
        port: u16,
        sender: EndpointServiceEventSender,
    ) -> bool {
        if self.senders.contains_key(&port) {
            return false;
        }
        self.senders.insert(port, sender);
        true
    }

    pub(in crate::node) fn is_registered(&self, port: u16) -> bool {
        self.senders.contains_key(&port)
    }

    pub(in crate::node) fn deliver(
        &self,
        messages: Vec<EndpointServiceDatagramDelivery>,
    ) -> Result<(), ()> {
        let mut by_port: HashMap<u16, Vec<EndpointServiceDatagramDelivery>> = HashMap::new();
        for message in messages {
            if self.senders.contains_key(&message.destination_port) {
                by_port
                    .entry(message.destination_port)
                    .or_default()
                    .push(message);
            }
        }
        for (port, messages) in by_port {
            if let Some(sender) = self.senders.get(&port) {
                sender.send(messages)?;
            }
        }
        Ok(())
    }
}
