use super::LinkTrait;
use core::fmt;
use futures_util::future::join_all;
use std::collections::HashSet;
use std::io;
use std::net::{IpAddr, SocketAddr};
use std::sync::Arc;
use tokio::net::UdpSocket;

pub struct UdpServer {
    addr: SocketAddr,
    socket: UdpSocket,
    data: [u8; 4096],
    clients: HashSet<SocketAddr>,
}
impl UdpServer {
    pub async fn new(addr: &str, port: u32) -> io::Result<UdpServer> {
        let addr: IpAddr = addr.parse().unwrap();
        let addr: SocketAddr = (addr, port as u16).into();
        let socket = UdpSocket::bind(&addr).await?;
        Ok(UdpServer {
            addr,
            socket,
            data: [0; 4096],
            clients: HashSet::new(),
        })
    }
}

#[async_trait::async_trait]
impl LinkTrait for UdpServer {
    async fn send(&mut self, msg: &[u8]) -> io::Result<()> {
        let send = self
            .clients
            .iter()
            .map(|addr| Box::pin(self.socket.send_to(msg, addr)));
        join_all(send).await;
        Ok(())
    }
    async fn recv(&mut self) -> io::Result<Arc<Vec<u8>>> {
        let (len, addr) = self.socket.recv_from(&mut self.data).await?;
        if self.clients.insert(addr) {
            log::info!("{} accept: {}", self, addr);
        }
        let data = Arc::new(Vec::from(&self.data[..len]));
        log::info!("{} recv: {:02X?}", self, data);
        Ok(data)
    }
}

impl fmt::Display for UdpServer {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.write_fmt(format_args!("UdpServer({:?})", self.addr))
    }
}