//! DNS-related protocol structures

use core::{
    fmt,
    mem::size_of,
    net::{IpAddr, Ipv4Addr},
};

use crate::util;
use alloc::{string::String, vec, vec::Vec};

// use super::IpAddr;

// TODO rewrite using NetValue

// "Wire" format

/// Raw DNS message header
#[derive(Clone, Copy)]
#[cfg_attr(feature = "bytemuck", derive(bytemuck::Pod, bytemuck::Zeroable))]
#[repr(C, packed)]
pub struct DnsHeader {
    /// Transaction ID
    pub xid: u16,
    /// Information about the message
    pub flags: u16,
    /// Question count
    pub qdcount: u16,
    /// Answer count
    pub ancount: u16,
    /// Nameserver record count
    pub nscount: u16,
    /// Additional record count
    pub arcount: u16,
}

/// DNS query/reply type
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[cfg_attr(feature = "bytemuck", derive(bytemuck::Pod, bytemuck::Zeroable))]
#[non_exhaustive]
#[repr(transparent)]
pub struct DnsType(u16);

/// DNS query/reply class
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[cfg_attr(feature = "bytemuck", derive(bytemuck::Pod, bytemuck::Zeroable))]
#[non_exhaustive]
#[repr(transparent)]
pub struct DnsClass(u16);

/// Describes whether a message is a query or a reply
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[cfg_attr(feature = "bytemuck", derive(bytemuck::Pod, bytemuck::Zeroable))]
#[non_exhaustive]
#[repr(transparent)]
pub struct DnsMethod(u16);

/// Describes the reply code of the message
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[cfg_attr(feature = "bytemuck", derive(bytemuck::Pod, bytemuck::Zeroable))]
#[non_exhaustive]
#[repr(transparent)]
pub struct DnsReplyCode(u16);

impl DnsType {
    /// A - Address (when used with IN class: IPv4)
    pub const A: Self = Self(0x0001);
    /// OPT - Option data
    pub const OPT: Self = Self(0x0029);
}

impl DnsClass {
    /// IN - Internet
    pub const IN: Self = Self(0x0001);
}

impl DnsMethod {
    /// The message is a query
    pub const QUERY: Self = Self(0);
    /// The message is a reply
    pub const REPLY: Self = Self(1);
}

impl DnsReplyCode {
    /// Successful status
    pub const NO_ERROR: Self = Self(0);
    /// Server could not recognize the message properly
    pub const FORMAT_ERROR: Self = Self(1);
    /// Server experienced an error
    pub const SERVER_ERROR: Self = Self(2);
    /// The name queried was invalid
    pub const NAME_ERROR: Self = Self(3);
    /// Feature is not implemented
    pub const NOT_IMPLEMENTED: Self = Self(4);
    /// The server refused to reply
    pub const REFUSED: Self = Self(5);
}

impl fmt::Display for DnsReplyCode {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let text = match *self {
            Self::NO_ERROR => "No error",
            Self::FORMAT_ERROR => "Format error",
            Self::SERVER_ERROR => "Server error",
            Self::NAME_ERROR => "Name error",
            Self::NOT_IMPLEMENTED => "Not implemented",
            Self::REFUSED => "Refused",
            _ => "<Unknown error>",
        };
        f.write_str(text)
    }
}

// "Local" format

/// Describes a single name query
#[derive(Debug)]
pub struct DnsQuery {
    /// Name being queried
    pub name: DnsName,
    /// Type
    pub ty: DnsType,
    /// Class
    pub class: DnsClass,
}

/// Describes an additional DNS option
#[derive(Debug)]
pub enum DnsOption {
    /// "Cookie" option to identify queries
    Cookie(u64),
}

/// Represents a DNS name: a sequence of dot-separated domain levels if Some or NULL if None
#[derive(Debug)]
pub struct DnsName(pub Option<String>);

/// Describes a DNS additional/nameserver/answer record data
#[derive(Debug)]
pub enum DnsRecordData {
    /// A-record: address
    A(Ipv4Addr),
    /// Option-record: additional options
    Option(DnsOption),
}

/// Describes a DNS record
#[derive(Debug)]
pub struct DnsRecord {
    /// DNS name this record is related to
    pub name: DnsName,
    /// Type of the record
    pub ty: DnsType,
    /// Class of the record
    pub class: DnsClass,
    /// TTL of the record (0 and ignored for Option records)
    pub ttl: u32,
    /// Data of the record
    pub rdata: DnsRecordData,
}

/// Describes a single DNS query/reply
#[derive(Debug)]
pub struct DnsMessage {
    /// Transaction ID
    pub xid: u16,
    /// Query or reply
    pub method: DnsMethod,
    /// Reply code
    pub reply_code: DnsReplyCode,
    /// Question section
    pub questions: Vec<DnsQuery>,
    /// Answer section
    pub answers: Vec<DnsRecord>,
    /// Additional section
    pub additional: Vec<DnsRecord>,
}

/// Helper trait for an abstract "buffer" which can be written to and extended
pub trait ExtensibleBuffer {
    /// Writes a slice of bytes to the buffer
    fn write_slice(&mut self, slice: &[u8]);
    /// Writes a byte to the buffer
    fn write_byte(&mut self, byte: u8);

    /// Writes value's bytes to the buffer
    #[cfg(feature = "bytemuck")]
    fn write_value<T: bytemuck::Pod>(&mut self, value: &T) {
        self.write_slice(bytemuck::bytes_of(value))
    }

    /// Writes value's bytes to the buffer
    #[cfg(not(feature = "bytemuck"))]
    fn write_value<T>(&mut self, value: &T) {
        self.write_slice(util::bytes_of(value))
    }
}

/// Helper trait for DNS message parsing/serialization
pub trait DnsSerialize {
    /// Serializes the message into the `output` buffer
    fn serialize<V: ExtensibleBuffer>(&self, output: &mut V);
}

impl DnsName {
    /// "NULL" DNS name: an empty name
    pub const NULL: Self = Self(None);
}

impl DnsMessage {
    /// Constructs a DNS query message
    pub fn query<S: Into<String>>(name: S, ty: DnsType, xid: u16, cookie: u64) -> Self {
        Self {
            xid,
            method: DnsMethod::QUERY,
            reply_code: DnsReplyCode::NO_ERROR,
            questions: vec![DnsQuery {
                name: DnsName(Some(name.into())),
                ty,
                class: DnsClass::IN,
            }],
            answers: vec![],
            additional: vec![DnsRecord {
                name: DnsName::NULL,
                ty: DnsType::OPT,
                class: DnsClass(1232),
                ttl: 0,
                rdata: DnsRecordData::Option(DnsOption::Cookie(cookie)),
            }],
        }
    }

    /// Constructs a DNS reply message
    pub fn reply(
        reply_code: DnsReplyCode,
        questions: Vec<DnsQuery>,
        answers: Vec<(String, IpAddr)>,
        xid: u16,
    ) -> Self {
        Self {
            xid,
            reply_code,
            method: DnsMethod::REPLY,
            questions,
            answers: answers
                .into_iter()
                .map(|(name, addr)| {
                    let name = DnsName(Some(name));
                    let (ty, rdata) = match addr {
                        IpAddr::V4(addr) => (DnsType::A, DnsRecordData::A(addr)),
                        _ => todo!(),
                    };

                    DnsRecord {
                        name,
                        ty,
                        class: DnsClass::IN,
                        ttl: 100,
                        rdata,
                    }
                })
                .collect(),
            additional: vec![],
        }
    }
}

impl ExtensibleBuffer for Vec<u8> {
    fn write_slice(&mut self, slice: &[u8]) {
        self.extend_from_slice(slice);
    }

    fn write_byte(&mut self, byte: u8) {
        self.push(byte);
    }
}

mod ser {
    use super::*;

    impl DnsSerialize for DnsName {
        fn serialize<V: ExtensibleBuffer>(&self, output: &mut V) {
            if let Some(name) = &self.0 {
                // Write question
                for word in name.split('.') {
                    if word.is_empty() {
                        continue;
                    }
                    output.write_byte(word.len().try_into().unwrap());
                    output.write_slice(word.as_bytes());
                }
            }
            // Zero-terminate
            output.write_byte(0);
        }
    }

    impl DnsSerialize for DnsOption {
        fn serialize<V: ExtensibleBuffer>(&self, output: &mut V) {
            match self {
                &Self::Cookie(value) => {
                    output.write_slice(&12u16.to_be_bytes()); // RDLENGTH
                    output.write_slice(&10u16.to_be_bytes()); // Option = COOKIE
                    output.write_slice(&8u16.to_be_bytes()); // Cookie length
                    output.write_slice(&value.to_be_bytes()); // Cookie data
                }
            }
        }
    }

    impl DnsSerialize for DnsRecordData {
        fn serialize<V: ExtensibleBuffer>(&self, output: &mut V) {
            match self {
                &Self::A(addr) => {
                    output.write_slice(&4u16.to_be_bytes());
                    output.write_slice(&u32::from(addr).to_be_bytes());
                }
                Self::Option(option) => option.serialize(output),
            }
        }
    }

    impl DnsSerialize for DnsRecord {
        fn serialize<V: ExtensibleBuffer>(&self, output: &mut V) {
            self.name.serialize(output);
            output.write_slice(&self.ty.0.to_be_bytes());
            output.write_slice(&self.class.0.to_be_bytes());
            output.write_slice(&self.ttl.to_be_bytes());
            self.rdata.serialize(output);
        }
    }

    impl DnsSerialize for DnsQuery {
        fn serialize<V: ExtensibleBuffer>(&self, output: &mut V) {
            self.name.serialize(output);
            output.write_slice(&self.ty.0.to_be_bytes());
            output.write_slice(&self.class.0.to_be_bytes());
        }
    }

    impl DnsSerialize for DnsMessage {
        fn serialize<V: ExtensibleBuffer>(&self, output: &mut V) {
            let qdcount: u16 = self.questions.len().try_into().unwrap();
            let ancount: u16 = self.answers.len().try_into().unwrap();
            let arcount: u16 = self.additional.len().try_into().unwrap();

            let header = DnsHeader {
                xid: self.xid.to_be(),
                flags: ((self.method.0 << 15) | (self.reply_code.0 & 0xF) | (1 << 8)).to_be(),
                qdcount: qdcount.to_be(),
                nscount: 0,
                ancount: ancount.to_be(),
                arcount: arcount.to_be(),
            };

            output.write_slice(util::bytes_of(&header));

            // Serialize data sections
            for question in self.questions.iter() {
                question.serialize(output);
            }
            for answer in self.answers.iter() {
                answer.serialize(output);
            }
            for additional in self.additional.iter() {
                additional.serialize(output);
            }
        }
    }
}

#[allow(missing_docs)]
mod de {
    use super::*;

    impl DnsName {
        pub fn parse(message: &[u8], input: &[u8]) -> Option<(Self, usize)> {
            let mut pos = 0;
            let mut value = String::new();

            loop {
                if input.len() < pos + 1 {
                    return None;
                }
                let len = input[pos];
                pos += 1;

                if len == 0 {
                    break;
                }

                match len & 0b11000000 {
                    0 => {
                        let len = len as usize;

                        if input.len() < pos + len {
                            return None;
                        }
                        let word = core::str::from_utf8(&input[pos..pos + len]).ok()?;
                        value.push_str(word);
                        value.push('.');
                        pos += len;
                    }
                    0b11000000 if len == 0b11000000 => {
                        // TODO pointer
                        if input.len() < pos + 1 {
                            return None;
                        }
                        let off = input[pos] as usize;
                        if input.len() < off {
                            return None;
                        }
                        let (rest, _) = Self::parse(message, &message[off..])?;
                        return Some((rest, pos + 1));
                    }
                    _ => return None,
                }
            }

            if value.is_empty() {
                Some((Self::NULL, pos))
            } else {
                Some((Self(Some(value)), pos))
            }
        }
    }

    impl DnsRecordData {
        pub fn parse(
            ty: DnsType,
            class: DnsClass,
            _message: &[u8],
            input: &[u8],
        ) -> Option<(Option<Self>, usize)> {
            if input.len() < size_of::<u16>() {
                return None;
            }
            let rdlength = u16::from_be_bytes([input[0], input[1]]) as usize;
            if input.len() < size_of::<u16>() + rdlength {
                return None;
            }

            let value = match (ty, class) {
                (DnsType::A, DnsClass::IN) if rdlength == 4 => {
                    let value = u32::from_be_bytes([input[2], input[3], input[4], input[5]]);
                    Some(Self::A(value.into()))
                }
                (DnsType::OPT, _) => {
                    todo!();
                }
                (_, _) => None,
            };

            Some((value, rdlength + size_of::<u16>()))
        }
    }

    impl DnsRecord {
        pub fn parse(message: &[u8], input: &[u8]) -> Option<(Option<Self>, usize)> {
            if input.len() < 1 + 4 * size_of::<u16>() {
                return None;
            }
            let (name, len) = DnsName::parse(message, input)?;
            if input.len() < len + 4 * size_of::<u16>() {
                return None;
            }
            let ty = DnsType(u16::from_be_bytes([input[len], input[len + 1]]));
            let class = DnsClass(u16::from_be_bytes([input[len + 2], input[len + 3]]));
            let ttl = u32::from_be_bytes([
                input[len + 4],
                input[len + 5],
                input[len + 6],
                input[len + 7],
            ]);
            let (rdata, rdata_len) = DnsRecordData::parse(ty, class, message, &input[len + 8..])?;

            let value = rdata.map(|v| Self {
                name,
                ty,
                class,
                ttl,
                rdata: v,
            });

            Some((value, rdata_len + len + 4 * size_of::<u16>()))
        }
    }

    impl DnsQuery {
        pub fn parse(message: &[u8], input: &[u8]) -> Option<(Self, usize)> {
            if input.len() < 1 + 2 * size_of::<u16>() {
                return None;
            }
            let (name, len) = DnsName::parse(message, input)?;
            if input.len() < len + 2 * size_of::<u16>() {
                return None;
            }
            let ty = DnsType(u16::from_be_bytes([input[len], input[len + 1]]));
            let class = DnsClass(u16::from_be_bytes([input[len + 2], input[len + 3]]));

            Some((Self { name, ty, class }, len + 2 * size_of::<u16>()))
        }
    }

    impl DnsMessage {
        pub fn parse(input: &[u8]) -> Option<Self> {
            if input.len() < size_of::<DnsHeader>() {
                return None;
            }
            let header: &DnsHeader = util::from_bytes(&input[..size_of::<DnsHeader>()]);

            let method = DnsMethod((u16::from_be(header.flags) >> 15) & 1);
            let qdcount = u16::from_be(header.qdcount);
            let ancount = u16::from_be(header.ancount);

            let reply_code = if method == DnsMethod::REPLY {
                let value = u16::from_be(header.flags) & 0xF;
                if value > 5 {
                    return None;
                }
                DnsReplyCode(value)
            } else {
                DnsReplyCode::NO_ERROR
            };

            // TODO don't ignore arcount/nscount

            let mut questions = vec![];
            let mut answers = vec![];

            let mut offset = size_of::<DnsHeader>();
            for _ in 0..qdcount {
                let (question, len) = DnsQuery::parse(input, &input[offset..])?;
                questions.push(question);
                offset += len;
            }
            for _ in 0..ancount {
                let (answer, len) = DnsRecord::parse(input, &input[offset..])?;
                if let Some(answer) = answer {
                    answers.push(answer);
                }
                offset += len;
            }

            Some(Self {
                xid: u16::from_be(header.xid),
                method,
                reply_code,
                questions,
                answers,
                additional: vec![],
            })
        }
    }
}
