use anyhow::{Context, Result};
use std::io;
use std::net::SocketAddr;
use std::sync::Arc;
use std::time::Duration;
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream, UdpSocket};
use tokio::task::{JoinError, JoinHandle, JoinSet};
use tokio_vsock::{VsockAddr, VsockStream};
use tracing::{info, warn};

const DEFAULT_DNS_PROXY_ADDR: &str = "127.0.0.1:53";
const SHUTTLE_DNS_PROXY_ADDR_ENV: &str = "SHUTTLE_DNS_PROXY_ADDR";

const MAX_DNS_MESSAGE_BYTES: usize = u16::MAX as usize;
const DNS_IO_TIMEOUT: Duration = Duration::from_secs(10);
const DNS_TCP_IDLE_TIMEOUT: Duration = Duration::from_secs(120);

// implements a proxy that sends dns requests to the spindle and
// lets the spindle resolve any queries, and streams the response back.
//
// we use this because the way spindle isolates QEMU VMs is unreliable
// to rely on. if we unblock the blackholed-routes for private nameservers
// like slirp one, we risk leaking internal DNS zones, even if the guest
// can't connect to them. there is also potentially DNS rebinding issues.
// and other slirp4netns quirks...
//
// this way we also get to filter the DNS queries very easily, so we can
// make sure we remove everything that would leak a host information.
pub struct DnsProxy {
    handles: Vec<JoinHandle<()>>,
}

#[derive(Clone)]
struct HostDnsClient {
    host_cid: u32,
    host_port: u32,
}

impl DnsProxy {
    pub async fn start(host_cid: u32, host_port: u32) -> Result<Option<Self>> {
        if host_port == 0 {
            return Ok(None);
        }

        let addr = std::env::var(SHUTTLE_DNS_PROXY_ADDR_ENV)
            .unwrap_or_else(|_| DEFAULT_DNS_PROXY_ADDR.to_owned());

        let udp = Arc::new(
            UdpSocket::bind(&addr)
                .await
                .with_context(|| format!("bind dns udp listener {addr}"))?,
        );

        let tcp = TcpListener::bind(&addr)
            .await
            .with_context(|| format!("bind dns tcp listener {addr}"))?;

        let host = HostDnsClient {
            host_cid,
            host_port,
        };

        let handles = vec![
            tokio::spawn(udp_loop(udp, host.clone())),
            tokio::spawn(tcp_loop(tcp, host)),
        ];

        info!(%addr, host_cid, host_port, "dns proxy ready");
        Ok(Some(Self { handles }))
    }
}

impl Drop for DnsProxy {
    fn drop(&mut self) {
        for handle in self.handles.drain(..) {
            handle.abort();
        }
    }
}

impl HostDnsClient {
    async fn query(&self, query: Vec<u8>) -> Result<Vec<u8>> {
        match self.query_once(&query).await {
            Ok(response) => Ok(response),
            Err(first_error) => self.query_once(&query).await.with_context(|| {
                format!("dns host query failed after retry; first error: {first_error:#}")
            }),
        }
    }

    async fn query_once(&self, query: &[u8]) -> Result<Vec<u8>> {
        let addr = VsockAddr::new(self.host_cid, self.host_port);

        let mut host = tokio::time::timeout(DNS_IO_TIMEOUT, VsockStream::connect(addr))
            .await
            .context("dns host connect timed out")?
            .with_context(|| {
                format!(
                    "dial host dns proxy cid={} port={}",
                    self.host_cid, self.host_port
                )
            })?;

        tokio::time::timeout(DNS_IO_TIMEOUT, async {
            write_dns_packet(&mut host, query)
                .await
                .context("write dns query to host")?;

            read_dns_packet(&mut host)
                .await
                .context("read dns response from host")?
                .context("host dns proxy closed without response")
        })
        .await
        .context("dns host query timed out")?
    }
}

async fn udp_loop(socket: Arc<UdpSocket>, host: HostDnsClient) {
    let mut buf = vec![0; MAX_DNS_MESSAGE_BYTES];
    let mut tasks = JoinSet::new();

    loop {
        tokio::select! {
            received = socket.recv_from(&mut buf) => match received {
                Ok((len, peer)) => {
                    let query = buf[..len].to_vec();
                    let socket = socket.clone();
                    let host = host.clone();

                    tasks.spawn(async move {
                        if let Err(error) = handle_udp_query(socket, peer, query, host).await {
                            warn!(%peer, %error, "dns udp query failed");
                        }
                    });
                }
                Err(error) => warn!(%error, "dns udp recv failed"),
            },

            Some(result) = tasks.join_next(), if !tasks.is_empty() => {
                log_dns_task_result(result);
            }
        }
    }
}

async fn handle_udp_query(
    socket: Arc<UdpSocket>,
    peer: SocketAddr,
    query: Vec<u8>,
    host: HostDnsClient,
) -> Result<()> {
    let response = host.query(query).await?;

    socket
        .send_to(&response, peer)
        .await
        .context("send dns udp response")?;

    Ok(())
}

async fn tcp_loop(listener: TcpListener, host: HostDnsClient) {
    let mut tasks = JoinSet::new();

    loop {
        tokio::select! {
            accepted = listener.accept() => match accepted {
                Ok((conn, peer)) => {
                    let host = host.clone();

                    tasks.spawn(async move {
                        if let Err(error) = handle_tcp_conn(conn, host).await {
                            warn!(%peer, %error, "dns tcp connection failed");
                        }
                    });
                }
                Err(error) => warn!(%error, "dns tcp accept failed"),
            },

            Some(result) = tasks.join_next(), if !tasks.is_empty() => {
                log_dns_task_result(result);
            }
        }
    }
}

async fn handle_tcp_conn(mut tcp: TcpStream, host: HostDnsClient) -> Result<()> {
    loop {
        let query = tokio::time::timeout(DNS_TCP_IDLE_TIMEOUT, read_dns_packet(&mut tcp))
            .await
            .context("dns tcp idle timeout")?
            .context("read dns tcp query")?;

        let Some(query) = query else {
            return Ok(());
        };

        let response = host.query(query).await?;

        write_dns_packet(&mut tcp, &response)
            .await
            .context("write dns tcp response")?;
    }
}

async fn read_dns_packet<R>(reader: &mut R) -> io::Result<Option<Vec<u8>>>
where
    R: AsyncRead + Unpin,
{
    let mut len_buf = [0; 2];

    match reader.read_exact(&mut len_buf).await {
        Ok(_) => {}
        Err(error) if error.kind() == io::ErrorKind::UnexpectedEof => return Ok(None),
        Err(error) => return Err(error),
    }

    let len = u16::from_be_bytes(len_buf) as usize;
    if len == 0 {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            "empty dns packet",
        ));
    }

    let mut packet = vec![0; len];
    reader.read_exact(&mut packet).await?;
    Ok(Some(packet))
}

async fn write_dns_packet<W>(writer: &mut W, packet: &[u8]) -> io::Result<()>
where
    W: AsyncWrite + Unpin,
{
    if packet.is_empty() || packet.len() > MAX_DNS_MESSAGE_BYTES {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            format!("invalid dns packet size {}", packet.len()),
        ));
    }

    writer
        .write_all(&(packet.len() as u16).to_be_bytes())
        .await?;
    writer.write_all(packet).await?;
    writer.flush().await
}

fn log_dns_task_result(result: Result<(), JoinError>) {
    if let Err(error) = result {
        warn!(%error, "dns proxy task failed");
    }
}