use anyhow::{bail, Context, Result};
use byteorder::{NetworkEndian, WriteBytesExt};
use rustls::{self, Session};
use std::io::{Cursor, Read, Write};
use std::net::{SocketAddr, TcpStream, UdpSocket};
use std::str::FromStr;
use std::sync::Arc;
use std::time::Duration;
use toluol::{DnsFlags, DnsOpcode, DnsMessage};

mod args;
use args::ConnectionType;

// TODO
// - better padding
// - display metainformation about query (args, server name, time etc)
// - update dependencies

fn main() -> Result<()> {
    let bufsize = 4096; // seems reasonable
    let args = args::parse_args().context("Failed to parse cli args")?;
    let data = prepare_query(&args, bufsize)?;

    let (answer, bytes_recvd) = match args.connection_type {
        ConnectionType::UDP => send_query_udp(&args, bufsize, &data),
        ConnectionType::TCP => send_query_tcp(&args, bufsize, &data),
        ConnectionType::TLS => send_query_tls(&args, &data)
    }?;

    let res = DnsMessage::parse(&mut Cursor::new(&answer))
        .context("Failed to parse answer")?;
    println!("Got an answer ({} bytes):\n{}", bytes_recvd, res);

    Ok(())
}

fn prepare_query(args: &args::Args, bufsize: u16) -> Result<Vec<u8>> {
    // see https://tools.ietf.org/html/rfc6840#section-5.9 for why the cd flag is set
    let flags = DnsFlags::new(false, false, true, false, true, true);
    let msg = DnsMessage::new_query(
        &args.url, args.qtype, DnsOpcode::QUERY, flags, true, args.use_dnssec, bufsize
    ).context("Failed to create query")?;
    msg.encode().context("Failed to encode query")
}

fn send_query_udp(args: &args::Args, bufsize: u16, data: &[u8]) -> Result<(Vec<u8>, u16)> {
    let socket = UdpSocket::bind("0.0.0.0:0").context("Failed to bind to UDP socket")?;
    let mut res = vec![0; bufsize as usize]; // the query sets this as max size

    socket.set_write_timeout(Some(Duration::new(2, 0)))
        .context("Failed to set socket write timeout")?;
    socket.set_read_timeout(Some(Duration::new(2, 0)))
        .context("Failed to set socket read timeout")?;

    socket.connect(format!("{}:53", args.nameserver))
        .context(format!("Failed to connect to {}:53 via UDP", args.nameserver))?;
    socket.send(&data).context("Failed to send data to nameserver")?;

    let bytes_recvd = socket.recv(&mut res).context("Error while receiving data via UDP socket")?;
    res.resize(bytes_recvd, 0);

    Ok((res, bytes_recvd as u16))
}

fn send_query_tcp(args: &args::Args, bufsize: u16, data: &[u8]) -> Result<(Vec<u8>, u16)> {
    let mut socket = TcpStream::connect(format!("{}:53", args.nameserver))
        .context(format!("Failed to connect to {}:53 via UDP", args.nameserver))?;
    let mut res = Vec::with_capacity(bufsize as usize); // the query sets this as max size

    socket.set_write_timeout(Some(Duration::new(2, 0)))
        .context("Failed to set TCP stream write timeout")?;
    socket.set_read_timeout(Some(Duration::new(2, 0)))
        .context("Failed to set TCP stream read timeout")?;

    let mut msg = Vec::with_capacity(data.len() + 2);
    msg.write_u16::<NetworkEndian>(data.len() as u16)?;
    msg.extend_from_slice(data);
    socket.write_all(&msg).context("Failed to write data to TCP stream")?;

    socket.read_to_end(&mut res).context("Failed to read from TCP stream")?;

    // again, RFC 1035, Section 4.2.2
    let bytes_recvd = ((res.remove(0) as u16) << 8) + (res.remove(0) as u16);
    if bytes_recvd != res.len() as u16 {
        bail!("Received {} bytes, but TCP message says {} were sent.", bytes_recvd, res.len())
    }

    Ok((res, bytes_recvd))
}

fn send_query_tls(args: &args::Args, data: &[u8]) -> Result<(Vec<u8>, u16)> {
    let mut config = rustls::ClientConfig::new();
    config.root_store.add_server_trust_anchors(&webpki_roots::TLS_SERVER_ROOTS);

    let hostname = args.hostname.as_ref().unwrap().as_str();
    let nameserver = webpki::DNSNameRef::try_from_ascii_str(hostname)
        .context("Invalid nameserver hostname")?;
    let mut session = rustls::ClientSession::new(&Arc::new(config), nameserver);

    let addr = SocketAddr::from_str(format!("{}:853", args.nameserver).as_str())
        .context("Invalid nameserver address")?;
    let mut socket = TcpStream::connect_timeout(&addr, Duration::new(2, 0))
        .context(format!("Failed to connect to {}:853, is the server configured to use DNS over TLS?", addr))?;
    let mut plaintext = Vec::new();

    socket.set_write_timeout(Some(Duration::new(2, 0)))
        .context("Failed to set TLS/TCP stream write timeout")?;
    socket.set_read_timeout(Some(Duration::new(2, 0)))
        .context("Failed to set TLS/TCP stream read timeout")?;

    let mut msg = Vec::with_capacity(data.len() + 2);
    msg.write_u16::<NetworkEndian>(data.len() as u16)?;
    msg.extend_from_slice(data);
    session.write_all(&msg).context("Failed to write to TLS socket")?;

    while plaintext.len() < 3 {
        if session.wants_write() {
            session.write_tls(&mut socket).context("Failed to write TLS packets to TCP stream")?;
        }

        if session.wants_read() {
            session.read_tls(&mut socket).context("Failed to read TLS packets from TCP stream")?;
            session.process_new_packets().context("Failed to process new TLS packets")?;
            session.read_to_end(&mut plaintext).context("Failed to read from TLS socket")?;
        }
    }

    session.send_close_notify();

    // remove first to bytes (see RFC 1035, Section 4.2.2)
    let bytes_recvd = ((plaintext.remove(0) as u16) << 8) + (plaintext.remove(0) as u16);
    if bytes_recvd != plaintext.len() as u16 {
        bail!("Received {} bytes, but TCP message says {} were sent.", bytes_recvd, plaintext.len())
    }

    Ok((plaintext, bytes_recvd))
}