#![feature(yggdrasil_os)]

use std::{
    collections::HashMap,
    os::{
        fd::AsRawFd,
        yggdrasil::io::{poll::PollChannel, timer::TimerFd},
    },
};
use std::{
    io,
    net::{IpAddr, Ipv4Addr, SocketAddr, UdpSocket},
    process::ExitCode,
    time::Duration,
};

use clap::{Parser, ValueEnum};
use yggdrasil_abi::net::dns::{
    DnsClass, DnsMessage, DnsMethod, DnsRecordData, DnsReplyCode, DnsSerialize, DnsType,
};

#[derive(thiserror::Error, Debug)]
enum Error {
    #[error("{0}")]
    IoError(#[from] io::Error),
    #[error("Timed out")]
    TimedOut,
}

#[derive(Parser)]
struct Arguments {
    #[clap(short, long)]
    server: Option<SocketAddr>,
    #[clap(short, long)]
    nameservers: Vec<IpAddr>,
    #[clap(short, long, value_enum)]
    record_type: Option<RecordType>,
    name: Option<String>,
}

#[derive(Clone, ValueEnum)]
enum RecordType {
    A,
}

#[derive(Debug, thiserror::Error)]
enum ServerError {
    #[error("{0}")]
    Error(#[from] Error),
    #[error("Invalid query")]
    InvalidQuery,
    #[error("Name does not exist")]
    NameDoesNotExist,
}

fn handle_query(
    nameservers: &[SocketAddr],
    a_cache: &mut HashMap<String, Ipv4Addr>,
    ty: DnsType,
    class: DnsClass,
    name: &str,
) -> Result<IpAddr, ServerError> {
    match (ty, class) {
        (DnsType::A, DnsClass::IN) => {
            if let Some(ip) = a_cache.get(name) {
                log::debug!("Cache hit: {} -> {}", name, ip);
                return Ok(IpAddr::V4(*ip));
            }

            log::debug!("Cache miss: {}", name);
            let result = perform_query(nameservers, name, ty)?;

            for answer in result.answers {
                let Some(name) = answer.name.0 else {
                    continue;
                };

                #[allow(clippy::single_match)]
                match (answer.ty, answer.class, answer.rdata) {
                    (DnsType::A, DnsClass::IN, DnsRecordData::A(ip)) => {
                        log::debug!("Resolved: {} -> {}", name, ip);
                        a_cache.insert(name, ip);
                    }
                    _ => (),
                }
            }

            a_cache
                .get(name)
                .copied()
                .map(IpAddr::V4)
                .ok_or(ServerError::NameDoesNotExist)
        }
        _ => Err(ServerError::InvalidQuery),
    }
}

fn run_server(addr: SocketAddr, nameservers: &[SocketAddr]) -> Result<(), Error> {
    // TODO cache TTL
    let mut a_cache = HashMap::new();
    let mut buffer = [0; 2048];
    let socket = UdpSocket::bind(addr)?;

    loop {
        let (len, remote) = socket.recv_from(&mut buffer)?;
        let Some(message) = DnsMessage::parse(&buffer[..len]) else {
            continue;
        };

        let mut answers = vec![];
        for (name, ty, class) in message
            .questions
            .iter()
            .filter_map(|q| q.name.0.as_deref().map(|name| (name, q.ty, q.class)))
        {
            match handle_query(nameservers, &mut a_cache, ty, class, name) {
                Ok(ip) => answers.push((name.to_owned(), ip)),
                Err(_) => {
                    continue;
                }
            }
        }

        let reply_code = if answers.len() == message.questions.len() {
            DnsReplyCode::NO_ERROR
        } else {
            DnsReplyCode::NAME_ERROR
        };

        let mut packet = vec![];
        let reply = DnsMessage::reply(reply_code, message.questions, answers, message.xid);
        reply.serialize(&mut packet);
        if let Err(error) = socket.send_to(&packet, remote) {
            log::error!("sendto: {}", error);
        }
    }
}

fn perform_query(nameservers: &[SocketAddr], name: &str, ty: DnsType) -> Result<DnsMessage, Error> {
    let mut buffer = [0; 2048];
    let socket = UdpSocket::bind("0.0.0.0:0")?;
    let mut timer = TimerFd::new(false, false)?;
    let mut poll = PollChannel::new()?;

    let timer_fd = timer.as_raw_fd();
    let socket_fd = socket.as_raw_fd();

    poll.add(socket_fd)?;
    poll.add(timer_fd)?;

    for _ in 0..5 {
        let mut packet = vec![];
        let xid = rand::random();
        let cookie = rand::random();
        let query = DnsMessage::query(name, ty, xid, cookie);
        query.serialize(&mut packet);

        for nameserver in nameservers.iter() {
            socket.send_to(&packet, nameserver)?;
        }
        timer.start(Duration::from_millis(500))?;

        loop {
            let (poll_fd, result) = poll.wait(None, true)?.unwrap();
            result?;

            match poll_fd {
                fd if fd == socket_fd => {
                    let (len, _) = socket.recv_from(&mut buffer)?;
                    let Some(message) = DnsMessage::parse(&buffer[..len]) else {
                        continue;
                    };
                    if message.xid != xid {
                        continue;
                    }
                    if message.method != DnsMethod::REPLY {
                        eprintln!("WARN: server sent a QUERY instead of REPLY");
                        continue;
                    }

                    return Ok(message);
                }
                fd if fd == timer_fd => {
                    break;
                }
                _ => unreachable!(),
            }
        }
    }

    Err(Error::TimedOut)
}

fn main() -> ExitCode {
    logsink::setup_logging(true);
    let mut args: Arguments = Arguments::parse();

    if let Some(server) = args.server {
        let nameservers: Vec<_> = args
            .nameservers
            .into_iter()
            .map(|ip| SocketAddr::new(ip, 53))
            .collect();

        if nameservers.is_empty() {
            eprintln!("-s requires one or more nameservers");
            return ExitCode::FAILURE;
        }

        match run_server(server, &nameservers) {
            Ok(()) => ExitCode::SUCCESS,
            Err(error) => {
                eprintln!("Error: {}", error);
                ExitCode::FAILURE
            }
        }
    } else if let Some(name) = args.name {
        if args.nameservers.is_empty() {
            args.nameservers
                .push(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)));
        }
        let nameservers: Vec<_> = args
            .nameservers
            .into_iter()
            .map(|ip| SocketAddr::new(ip, 53))
            .collect();

        let ty = match args.record_type.unwrap_or(RecordType::A) {
            RecordType::A => DnsType::A,
        };

        let message = match perform_query(&nameservers, &name, ty) {
            Ok(message) => message,
            Err(error) => {
                eprintln!("Error: {}", error);
                return ExitCode::FAILURE;
            }
        };

        if message.reply_code == DnsReplyCode::NO_ERROR {
            for answer in message.answers.iter() {
                let name = answer.name.0.as_deref().unwrap_or("<NULL>");
                #[allow(clippy::single_match)]
                match (answer.ty, answer.class, &answer.rdata) {
                    (DnsType::A, DnsClass::IN, DnsRecordData::A(address)) => {
                        println!("{}: {}", name, address);
                    }
                    _ => (),
                }
            }

            ExitCode::SUCCESS
        } else {
            eprintln!("Server returned error: {}", message.reply_code);
            ExitCode::FAILURE
        }
    } else {
        eprintln!("Either -s LISTEN-ADDRESS or NAME must be specified");
        ExitCode::FAILURE
    }
}
