use std::convert::TryInto;

/// Taken from RFC 6895, 2.2. OpCode Assignment
/// Currently, DNS OpCodes are assigned as follows:
/// OpCode Name                                Reference
///
/// 0      Query                               [RFC1035]
/// 1      IQuery (Inverse Query, OBSOLETE)    [RFC3425]
/// 2      Status                              [RFC1035]
/// 3      Unassigned
/// 4      Notify                              [RFC1996]
/// 5      Update                              [RFC2136]
/// 6-15   Unassigned
#[derive(Debug, PartialEq, Copy, Clone)]
pub enum DNSOpCode {
    Query = 0,
    IQuery = 1, // obsolete
    Status = 2,
    Notify = 4,
    Update = 5,
}

impl TryFrom<u8> for DNSOpCode {
    type Error = DNSParseError;

    fn try_from(value: u8) -> Result<Self, Self::Error> {
        match value {
            0 => Ok(DNSOpCode::Query),
            1 => Ok(DNSOpCode::IQuery),
            2 => Ok(DNSOpCode::Status),
            4 => Ok(DNSOpCode::Notify),
            5 => Ok(DNSOpCode::Update),
            _ => Err(DNSParseError::DNSOpCodeInvalid),
        }
    }
}

/// Taken from RFC 6895, 2.3. RCODE Assignment
///
/// RCODE   Name    Description                        Reference
///
/// 0    NoError   No Error                           [RFC1035]
/// 1    FormErr   Format Error                       [RFC1035]
/// 2    ServFail  Server Failure                     [RFC1035]
/// 3    NXDomain  Non-Existent Domain                [RFC1035]
/// 4    NotImp    Not Implemented                    [RFC1035]
/// 5    Refused   Query Refused                      [RFC1035]
/// 6    YXDomain  Name Exists when it should not     [RFC2136]
/// 7    YXRRSet   RR Set Exists when it should not   [RFC2136]
/// 8    NXRRSet   RR Set that should exist does not  [RFC2136]
/// 9    NotAuth   Server Not Authoritative for zone  [RFC2136]
/// 9    NotAuth   Not Authorized                     [RFC2845]
/// 10   NotZone   Name not contained in zone         [RFC2136]
///
/// 11 - 15 Unassigned
///
/// 16   BADVERS   Bad OPT Version                    [RFC6891]
/// 16   BADSIG    TSIG Signature Failure             [RFC2845]
/// 17   BADKEY    Key not recognized                 [RFC2845]
/// 18   BADTIME   Signature out of time window       [RFC2845]
/// 19   BADMODE   Bad TKEY Mode                      [RFC2930]
/// 20   BADNAME   Duplicate key name                 [RFC2930]
/// 21   BADALG    Algorithm not supported            [RFC2930]
/// 22   BADTRUNC  Bad Truncation                     [RFC4635]
///
/// 23 - 3,840
/// 0x0017 - 0x0F00     Unassigned
///
/// 3,841 - 4,095
/// 0x0F01 - 0x0FFF     Reserved for Private Use
///
/// 4,096 - 65,534
/// 0x1000 - 0xFFFE     Unassigned
///
/// 65,535
/// 0xFFFF              Reserved; can only be allocated by Standards Action.
#[derive(Debug)]
pub enum DNSRCode {
    NoError = 0,
    FormErr = 1,
    ServFail = 2,
    NXDomain = 3,
    NotImp = 4,
    Refused = 5,
    YXDomain = 6,
    YXRRSet = 7,
    NXRRSet = 8,
    NotAuth = 9,
    NotZone = 10,
}

impl TryFrom<u8> for DNSRCode {
    type Error = DNSParseError;

    fn try_from(value: u8) -> Result<Self, Self::Error> {
        match value {
            0 => Ok(DNSRCode::NoError),
            1 => Ok(DNSRCode::FormErr),
            2 => Ok(DNSRCode::ServFail),
            3 => Ok(DNSRCode::NXDomain),
            4 => Ok(DNSRCode::NotImp),
            5 => Ok(DNSRCode::Refused),
            6 => Ok(DNSRCode::YXDomain),
            7 => Ok(DNSRCode::YXRRSet),
            8 => Ok(DNSRCode::NXRRSet),
            9 => Ok(DNSRCode::NotAuth),
            10 => Ok(DNSRCode::NotZone),
            _ => Err(DNSParseError::DNSRCodeInvalid),
        }
    }
}

#[derive(Debug)]
pub struct DNSHeader {
    /// used by the requester to match up replies to outstanding queries
    pub id: u16,
    /// specifies whether this message is a query (false), or a response (true)
    pub qr: bool,
    pub opcode: DNSOpCode,
    /// specifies that the responding name server is an authority for the domain name in question section
    pub aa: bool,
    /// specifies that this message was truncated due to length greater than that permitted on the transmission channel
    pub tc: bool,
    /// it directs the name server to pursue the query recursively
    pub rd: bool,
    /// denotes whether recursive query support is available in the name server
    pub ra: bool,
    /// TODO: add documuentation about this flag
    pub ad: bool,
    /// TODO: add documuentation about this flag
    pub cd: bool,
    pub rcode: DNSRCode,
    /// TODO: add documuentation about this count
    pub qd_zo_count: u16,
    /// TODO: add documuentation about this count
    pub an_pr_count: u16,
    /// TODO: add documuentation about this count
    pub ns_up_count: u16,
    /// TODO: add documuentation about this count
    pub ar_count: u16,
}

#[derive(Debug)]
pub struct DNSQuery {
    pub hdr: DNSHeader,
    pub name: String,
    pub qclass: u16,
    pub qtype: u16,
}

impl DNSHeader {
    const QR_MASK: u8 = 0b10000000;
    const OPCODE_MASK: u8 = 0b01111000;
    const AA_MASK: u8 = 0b00000100;
    const TC_MASK: u8 = 0b00000010;
    const RD_MASK: u8 = 0b00000001;

    const RA_MASK: u8 = 0b10000000;
    const AD_MASK: u8 = 0b00100000;
    const CD_MASK: u8 = 0b00010000;
    const RCODE_MASK: u8 = 0b00001111;

    const OPCODE_OFFSET: u8 = 3;

    pub fn from_udp_datagram(datagram: &[u8]) -> Result<Self, DNSParseError> {
        if datagram.len() < 12 {
            return Err(DNSParseError::DatagramTooShort);
        }

        let id = u16::from_be_bytes((&datagram[..2]).try_into().unwrap());

        let qr = (datagram[2] & Self::QR_MASK) != 0;
        let opcode = DNSOpCode::try_from((datagram[2] & Self::OPCODE_MASK) >> Self::OPCODE_OFFSET)?;
        let aa = (datagram[2] & Self::AA_MASK) != 0;
        let tc = (datagram[2] & Self::TC_MASK) != 0;
        let rd = (datagram[2] & Self::RD_MASK) != 0;

        let ra = (datagram[3] & Self::RA_MASK) != 0;
        let ad = (datagram[3] & Self::AD_MASK) != 0;
        let cd = (datagram[3] & Self::CD_MASK) != 0;
        let rcode = DNSRCode::try_from(datagram[3] & Self::RCODE_MASK)?;

        let qd_zo_count = u16::from_be_bytes((datagram[4..6]).try_into().unwrap());
        let an_pr_count = u16::from_be_bytes((datagram[6..8]).try_into().unwrap());
        let ns_up_count = u16::from_be_bytes((datagram[8..10]).try_into().unwrap());
        let ar_count = u16::from_be_bytes((datagram[10..12]).try_into().unwrap());

        Ok(DNSHeader {
            id,
            qr,
            opcode,
            aa,
            tc,
            rd,
            ra,
            ad,
            cd,
            rcode,
            qd_zo_count,
            an_pr_count,
            ns_up_count,
            ar_count,
        })
    }
}

#[derive(Debug, PartialEq)]
pub enum DNSParseError {
    DatagramTooShort,
    DNSOpCodeInvalid,
    DNSRCodeInvalid,
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn parse_dns_header_too_short() {
        // minimal header with 1 byte missing
        let dns_query = [
            0x13, 0x37, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00,
        ];
        let parse_result = DNSHeader::from_udp_datagram(&dns_query);
        assert_eq!(parse_result.err().unwrap(), DNSParseError::DatagramTooShort);
    }

    #[test]
    fn parse_dns_header_opcode_invalid() {
        let mut dns_query = [
            0xff, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
        ];
        let invalid_opcodes = [3, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15];

        for opcode in invalid_opcodes {
            dns_query[2] = opcode << 3;
            let parse_result = DNSHeader::from_udp_datagram(&dns_query);
            assert_eq!(
                parse_result.err(),
                Some(DNSParseError::DNSOpCodeInvalid),
                "query: {:02x?}, opcode: {}",
                dns_query,
                opcode
            );
        }
    }

    #[test]
    fn parse_dns_header_opcode_valid() -> Result<(), DNSParseError> {
        let mut dns_query = [
            0xff, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
        ];
        let valid_opcodes = [0, 1, 2, 4, 5];
        let valid_parsed_opcodes = [
            DNSOpCode::Query,
            DNSOpCode::IQuery,
            DNSOpCode::Status,
            DNSOpCode::Notify,
            DNSOpCode::Update,
        ];

        for (idx, opcode) in valid_opcodes.iter().enumerate() {
            dns_query[2] = opcode << 3;
            let parse_result = DNSHeader::from_udp_datagram(&dns_query);
            assert_eq!(
                parse_result?.opcode, valid_parsed_opcodes[idx],
                "query: {:02x?}, opcode: {}",
                dns_query, opcode
            );
        }
        Ok(())
    }
}