use dnssector::DNS_MAX_COMPRESSED_SIZE;
use std::{
    io::{self, Read, Write},
    net::{SocketAddr, TcpStream, UdpSocket},
    time::Duration,
};

pub fn query_raw_udp(
    local: &SocketAddr,
    upstream: &SocketAddr,
    packet: &[u8],
    timeout: Duration,
) -> io::Result<Vec<u8>> {
    let socket = UdpSocket::bind(local)?;
    let _ = socket.set_read_timeout(Some(timeout));
    socket.connect(upstream)?;
    socket.send(packet)?;
    let mut response = vec![0; DNS_MAX_COMPRESSED_SIZE];
    let response_len = socket.recv(&mut response)?;
    response.truncate(response_len);
    Ok(response)
}

pub fn query_raw_tcp(
    upstream: &SocketAddr,
    packet: &[u8],
    timeout: Duration,
) -> io::Result<Vec<u8>> {
    let mut stream = TcpStream::connect_timeout(upstream, timeout)?;
    let _ = stream.set_read_timeout(Some(timeout));
    let _ = stream.set_write_timeout(Some(timeout));
    let _ = stream.set_nodelay(true);
    let query_len = packet.len();
    let mut tcp_query = Vec::with_capacity(2 + query_len);
    tcp_query.push((query_len >> 8) as u8);
    tcp_query.push(query_len as u8);
    tcp_query.extend_from_slice(packet);
    stream.write_all(&tcp_query)?;
    let mut response_len_bytes = [0u8; 2];
    stream.read_exact(&mut response_len_bytes)?;
    let response_len = ((response_len_bytes[0] as usize) << 8) | (response_len_bytes[1] as usize);
    let mut response = vec![0; response_len];
    stream.read_exact(&mut response)?;
    Ok(response)
}