#![warn(clippy::pedantic)]
#![warn(clippy::nursery)]
#![warn(clippy::unwrap_used)]
#![warn(clippy::todo)]
#![warn(clippy::panic)]

mod config;
mod dns;
mod ext;
mod logging;
mod macros;

use std::{io::Cursor, net::SocketAddr};

use anyhow::Result;
use dotenv::dotenv;
use tokio::net::UdpSocket;
use tracing::{Instrument, Level};

use crate::{
    config::ConfigData,
    dns::{Answer, CLASS_IN, Message, Question, RECORD_TYPE_A, RESPONSE_CODE_NXDOMAIN},
};

#[tokio::main]
async fn main() -> Result<()> {
    dotenv().ok();
    logging::setup();
    let config = config::load()?;

    tracing::info!("start server on port {}", config.port);
    let sock = UdpSocket::bind(("0.0.0.0", config.port)).await?;
    let mut buffer = [0u8; 512];

    loop {
        let (size, addr) = sock.recv_from(&mut buffer).await?;
        let mut cursor = Cursor::new(&buffer[..size]);
        let result = handle_datagram(&config, &sock, &addr, &mut cursor).await;
        if let Err(err) = result {
            tracing::error!("failed to handle datagram: {err}");
        }
    }
}

async fn handle_datagram(
    config: &config::ConfigData,
    sock: &UdpSocket,
    addr: &SocketAddr,
    cursor: &mut Cursor<&[u8]>,
) -> Result<()> {
    let message = Message::parse(cursor)?;
    let span = tracing::span!(
        Level::INFO,
        "handle_datagram",
        transaction_id = %txid!(message.header.transaction_id),
        client = %addr,
    );

    let result: Result<()> = async {
        tracing::info!("received message");

        if message.header.query_response {
            tracing::error!("message is response, drop");
            return Ok(());
        }
        if message.header.op_code != 0 {
            tracing::error!("op code is not 0, drop");
            return Ok(());
        }

        let mut response = message.prepare_response();
        for question in message.questions {
            if question.class != CLASS_IN {
                tracing::error!(
                    "only IN class records are supported yet, requested: {}",
                    question.class
                );
                return Ok(());
            }

            match question.qtype {
                RECORD_TYPE_A => {
                    handle_type_a_question(config, &question, &mut response);
                }
                qt => {
                    tracing::error!("unsupported question type, requested: {qt}");
                    return Ok(());
                }
            }
        }

        let response_bytes = response.to_bytes()?;
        sock.send_to(&response_bytes, addr).await?;

        Ok(())
    }
    .instrument(span)
    .await;

    result
}

fn handle_type_a_question(config: &ConfigData, question: &Question, response: &mut Message) {
    let address = config.get_ipv4_address_by_domain(&question.domain);
    match address {
        Some(address) => {
            let answer = Answer {
                domain: question.domain.clone(),
                record_type: RECORD_TYPE_A,
                class: CLASS_IN,
                ttl: 1,
                data: address.to_vec(),
            };

            response.answers.push(answer);
            response.header.number_of_answers += 1;
        }
        None => {
            response.header.response_code = RESPONSE_CODE_NXDOMAIN;
        }
    }
}