use std::pin::Pin;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::Arc;
use std::task::{Context, Poll};

use bytes::Bytes;
use tokio::io::ReadBuf;
use tokio::net::UdpSocket;
use tokio::sync::mpsc::{self, Receiver, Sender};
use tokio::task::JoinHandle;

use crate::async_stream::{
    AsyncFlushMessage, AsyncMessageStream, AsyncPing, AsyncReadMessage, AsyncShutdownMessage,
    AsyncWriteMessage,
};

pub struct UdpMultiSocket {
    join_handles: Vec<JoinHandle<()>>,
    next_index: AtomicUsize,
    senders: Vec<Sender<Bytes>>,
    receiver: Receiver<Bytes>,
    notify_shutdown: Arc<AtomicBool>,
}

impl UdpMultiSocket {
    pub fn new(socket: UdpSocket, sender_count: usize, receiver_count: usize) -> Self {
        let socket = Arc::new(socket);
        let shutdown_flag = Arc::new(AtomicBool::new(false));
        let (agg_tx, agg_rx) = mpsc::channel(100);
        let mut join_handles = Vec::new();
        let mut senders = Vec::new();

        // Spawn sender tasks.
        for _ in 0..sender_count {
            let (tx, mut rx) = mpsc::channel::<Bytes>(100);
            senders.push(tx);
            let sock_clone = socket.clone();
            let shutdown_flag = shutdown_flag.clone();
            join_handles.push(tokio::spawn(async move {
                loop {
                    if shutdown_flag.load(Ordering::Relaxed) {
                        break;
                    }
                    match rx.recv().await {
                        Some(message) => {
                            if let Err(e) = sock_clone.send(&message).await {
                                eprintln!("send error: {:?}", e);
                                break;
                            }
                        }
                        None => break,
                    }
                }
            }));
        }

        // Spawn receiver tasks.
        for _ in 0..receiver_count {
            let sock_clone = socket.clone();
            let agg_tx = agg_tx.clone();
            let shutdown_flag = shutdown_flag.clone();
            join_handles.push(tokio::spawn(async move {
                let mut buf = [0u8; 65535];
                loop {
                    if shutdown_flag.load(Ordering::Relaxed) {
                        break;
                    }
                    match sock_clone.recv(&mut buf).await {
                        Ok(n) => {
                            let message = Bytes::copy_from_slice(&buf[..n]);
                            if agg_tx.send(message).await.is_err() {
                                break;
                            }
                        }
                        Err(e) => {
                            eprintln!("recv error: {:?}", e);
                            break;
                        }
                    }
                }
            }));
        }
        drop(agg_tx);

        UdpMultiSocket {
            join_handles,
            next_index: AtomicUsize::new(0),
            senders,
            receiver: agg_rx,
            notify_shutdown: shutdown_flag,
        }
    }
}

impl AsyncReadMessage for UdpMultiSocket {
    fn poll_read_message(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<std::io::Result<()>> {
        let this = self.get_mut();
        match Pin::new(&mut this.receiver).poll_recv(cx) {
            Poll::Ready(Some(message)) => {
                if message.len() > buf.remaining() {
                    return Poll::Ready(Err(std::io::Error::new(
                        std::io::ErrorKind::Other,
                        "buffer too small",
                    )));
                }
                buf.put_slice(&message);
                Poll::Ready(Ok(()))
            }
            Poll::Ready(None) => Poll::Ready(Err(std::io::Error::new(
                std::io::ErrorKind::UnexpectedEof,
                "channel closed",
            ))),
            Poll::Pending => Poll::Pending,
        }
    }
}

impl AsyncWriteMessage for UdpMultiSocket {
    fn poll_write_message(
        self: Pin<&mut Self>,
        _cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<std::io::Result<()>> {
        let this = self.get_mut();
        let index = this.next_index.fetch_add(1, Ordering::Relaxed) % this.senders.len();
        let msg = Bytes::copy_from_slice(buf);
        match this.senders[index].try_send(msg) {
            Ok(()) => Poll::Ready(Ok(())),
            Err(mpsc::error::TrySendError::Full(_)) => Poll::Pending,
            Err(e) => Poll::Ready(Err(std::io::Error::new(
                std::io::ErrorKind::Other,
                e.to_string(),
            ))),
        }
    }
}

impl AsyncFlushMessage for UdpMultiSocket {
    fn poll_flush_message(
        self: Pin<&mut Self>,
        _cx: &mut Context<'_>,
    ) -> Poll<std::io::Result<()>> {
        Poll::Ready(Ok(()))
    }
}

impl AsyncShutdownMessage for UdpMultiSocket {
    fn poll_shutdown_message(
        self: Pin<&mut Self>,
        _cx: &mut Context<'_>,
    ) -> Poll<std::io::Result<()>> {
        let this = self.get_mut();
        this.notify_shutdown.store(true, Ordering::Relaxed);
        this.senders.clear();
        for handle in this.join_handles.drain(..) {
            handle.abort();
        }
        Poll::Ready(Ok(()))
    }
}

impl AsyncPing for UdpMultiSocket {
    fn supports_ping(&self) -> bool {
        false
    }

    fn poll_write_ping(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<bool>> {
        unimplemented!();
    }
}

impl AsyncMessageStream for UdpMultiSocket {}
