use dnsparse::{Answer, Header, HeaderKind, Message, OpCode, QueryClass, QueryKind, ResponseCode};
use embassy_net::{
    Stack,
    udp::{PacketMetadata, UdpSocket},
};
use embassy_time::Timer;
// use log::{info, warn};
use defmt::*;
// use pico_website::{unwrap, unwrap_opt};

#[embassy_executor::task(pool_size = 1)]
pub async fn dns_server(stack: Stack<'static>) {
    let mut rx_buffer = [0; 4096];
    let mut tx_buffer = [0; 4096];
    let mut rx_meta = [PacketMetadata::EMPTY; 16];
    let mut tx_meta = [PacketMetadata::EMPTY; 16];
    let mut buf = [0; 4096];
    let mut res_buf = [0; 2048];
    loop {
        let mut socket = UdpSocket::new(
            stack,
            &mut rx_meta,
            &mut rx_buffer,
            &mut tx_meta,
            &mut tx_buffer,
        );
        unwrap!(socket.bind(53));

        info!("Starting DNS server");
        loop {
            Timer::after_secs(0).await;
            let (n, meta) = unwrap!(socket.recv_from(&mut buf).await);

            let msg = unwrap!(buf.get_mut(..n));
            let msg = match dnsparse::Message::parse(msg) {
                Ok(msg) => msg,
                Err(e) => {
                    warn!(
                        "Dns: Error while parsing DNS message : {}",
                        Display2Format(&e)
                    );
                    continue;
                }
            };

            if msg.header().opcode() != OpCode::Query {
                info!(
                    "Dns: Received unknown dns opcode ({:?}), ignoring",
                    Debug2Format(&msg.header().opcode())
                );
                continue;
            }
            let mut res = Message::builder(&mut res_buf)
                .header(
                    Header::builder()
                        .id(msg.header().id())
                        .kind(HeaderKind::Response)
                        .recursion_available(false)
                        .recursion_desired(msg.header().recursion_desired())
                        .response_code(ResponseCode::NoError)
                        .build(),
                )
                .build();
            for q in msg.questions() {
                match q.kind() {
                    QueryKind::A => {
                        if q.name() == "pico.wifi" || q.name() == "www.pico.wifi" {
                            res.add_question(&q);
                            res.add_answer(&Answer {
                                name: q.name().clone(),
                                kind: QueryKind::A,
                                class: QueryClass::IN,
                                ttl: 600,
                                rdata: &[192, 254, 0, 2],
                            });
                            info!("Dns: Giving {}", Display2Format(&q.name()));
                        } else {
                            info!("Dns: Unknown uri, ignoring ({})", Display2Format(&q.name()));
                        }
                    }
                    _ => {
                        continue;
                    }
                };
            }

            if let Err(e) = socket.send_to(res.as_bytes(), meta).await {
                warn!("Dns: Error while sending dns response : {:?}", e);
                break;
            }
        }
    }
}