use std::net::SocketAddr;

use anyhow::Result;
use async_net::UdpSocket;
use bytes::BytesMut;
use tracing::{self, trace, warn};

use crate::{lookup, runner};

/// Listens for UDP requests until an accept fails or the channel sender is closed.
pub async fn listen_udp(
    udp_recv: &mut UdpSocket,
    server_query_tx: &mut async_channel::Sender<runner::RequestMsg>,
) -> Result<()> {
    let mut packet_buffer = BytesMut::with_capacity(4096);
    loop {
        // Ensure that the buffer has a SIZE suitable for socket.recv_from().
        // If we just leave it with the CAPACITY then it drops data.
        packet_buffer.resize(packet_buffer.capacity(), 0);

        let (recvsize, request_source) = udp_recv.recv_from(&mut packet_buffer).await?;

        // Got a request from somewhere, shorten to actual size received (doesnt affect malloc)
        packet_buffer.truncate(recvsize);

        trace!(
            "Queueing raw UDP request from {:?} ({}b): {:02X?}",
            request_source,
            packet_buffer.len(),
            &packet_buffer[..]
        );
        let msg = runner::RequestMsg {
            src: request_source,
            data: runner::RequestData::Udp(packet_buffer.clone()),
        };

        // Forward request info to be processed and then replied
        server_query_tx.send(msg).await?;
    }
}

pub async fn handle_udp_request(
    lookup: &mut lookup::Lookup,
    request_source: SocketAddr,
    mut buf: BytesMut,
    udp_sock: &UdpSocket,
) {
    match lookup.handle_query(&mut buf).await {
        Ok(()) => {
            // Send the response back to the original UDP requestor.
            trace!("Raw response to {:?} ({}b): {:02X?}", request_source, buf.len(), &buf[..]);
            if let Err(e) = udp_sock.send_to(&mut buf, &request_source).await {
                warn!(
                    "Failed to send UDP DNS response to client={:?}: {:?}",
                    request_source, e
                );
            }
        }
        Err(ioerr) => {
            warn!(
                "Failed to handle request from client={:?}: {:02X?} ({:?})",
                request_source,
                &buf[..],
                ioerr
            );
        }
    }
}