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, PartialEq, Copy, Clone)]
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, PartialEq, Clone, Copy)]
pub enum DNSMessageType {
    Query,
    Response,
}

impl From<bool> for DNSMessageType {
    fn from(value: bool) -> DNSMessageType {
        match value {
            false => Self::Query,
            true => Self::Response,
        }
    }
}

#[derive(Debug, Clone, PartialEq)]
pub struct DNSHeader {
    /// used by the requester to match up replies to outstanding queries
    pub id: u16,
    /// specifies whether this message is a query or a response
    pub message_type: DNSMessageType,
    pub opcode: DNSOpCode,
    /// specifies that the responding name server is an authority for the domain name in question section
    pub authorative_answer: bool,
    /// specifies that this message was truncated due to length greater than that permitted on the transmission channel
    pub truncated: bool,
    /// it directs the name server to pursue the query recursively
    pub recursion_desired: bool,
    /// denotes whether recursive query support is available in the name server
    pub recursion_available: bool,
    /// TODO: add documuentation about this flag
    pub authentic_data: bool,
    /// TODO: add documuentation about this flag
    pub checking_disabled: bool,
    pub response_code: DNSRCode,
    /// TODO: add documuentation about this count
    pub query_count: u16,
    /// TODO: add documuentation about this count
    pub answer_count: u16,
    /// TODO: add documuentation about this count
    pub name_server_count: u16,
    /// TODO: add documuentation about this count
    pub additional_count: u16,
}

#[derive(Debug, Clone)]
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 message_type = DNSMessageType::from(datagram[2] & Self::QR_MASK != 0);
        let opcode = DNSOpCode::try_from((datagram[2] & Self::OPCODE_MASK) >> Self::OPCODE_OFFSET)?;
        let authorative_answer = (datagram[2] & Self::AA_MASK) != 0;
        let truncated = (datagram[2] & Self::TC_MASK) != 0;
        let recursion_desired = (datagram[2] & Self::RD_MASK) != 0;

        let recursion_available = (datagram[2] & Self::RA_MASK) != 0;
        let authentic_data = (datagram[2] & Self::AD_MASK) != 0;
        let checking_disabled = (datagram[2] & Self::CD_MASK) != 0;
        let response_code = DNSRCode::try_from(datagram[3] & Self::RCODE_MASK)?;

        let query_count = u16::from_be_bytes((datagram[4..6]).try_into().unwrap());
        let answer_count = u16::from_be_bytes((datagram[6..8]).try_into().unwrap());
        let name_server_count = u16::from_be_bytes((datagram[8..10]).try_into().unwrap());
        let additional_count = u16::from_be_bytes((datagram[10..12]).try_into().unwrap());

        Ok(DNSHeader {
            id,
            message_type,
            opcode,
            authorative_answer,
            truncated,
            recursion_desired,
            recursion_available,
            authentic_data,
            checking_disabled,
            response_code,
            query_count,
            answer_count,
            name_server_count,
            additional_count,
        })
    }
}

#[derive(Debug, Clone, 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, DNSOpCode::Query),
            (1, DNSOpCode::IQuery),
            (2, DNSOpCode::Status),
            (4, DNSOpCode::Notify),
            (5, DNSOpCode::Update),
        ];

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

    #[test]
    fn parse_dns_header_message_type_query() -> Result<(), DNSParseError> {
        let dns_query = [
            0xff, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
        ];
        let parse_result = DNSHeader::from_udp_datagram(&dns_query);
        assert_eq!(parse_result?.message_type, DNSMessageType::Query);
        Ok(())
    }

    #[test]
    fn parse_dns_header_message_type_response() -> Result<(), DNSParseError> {
        let dns_query = [
            0xff, 0x00, 0x80, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
        ];
        let parse_result = DNSHeader::from_udp_datagram(&dns_query);
        assert_eq!(parse_result?.message_type, DNSMessageType::Response);
        Ok(())
    }

    #[test]
    fn parse_dns_header_response_code_invalid() {
        let mut dns_query = [
            0xff, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
        ];
        let invalid_rcodes = [11, 12, 13, 14, 15];

        for rcode in invalid_rcodes {
            dns_query[3] = (rcode as u8) << 0;

            let parse_result = DNSHeader::from_udp_datagram(&dns_query);

            assert_eq!(
                parse_result,
                Err(DNSParseError::DNSRCodeInvalid),
                "query: {:02x?}, rcode: {}",
                dns_query,
                rcode,
            );
        }
    }

    #[test]
    fn parse_dns_header_response_code_valid() -> Result<(), DNSParseError> {
        let mut dns_query = [
            0xff, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
        ];

        // for now we only support RCodes 0 to 10 (inclusive)
        let valid_rcodes = [
            (0, DNSRCode::NoError),
            (1, DNSRCode::FormErr),
            (2, DNSRCode::ServFail),
            (3, DNSRCode::NXDomain),
            (4, DNSRCode::NotImp),
            (5, DNSRCode::Refused),
            (6, DNSRCode::YXDomain),
            (7, DNSRCode::YXRRSet),
            (8, DNSRCode::NXRRSet),
            (9, DNSRCode::NotAuth),
            (10, DNSRCode::NotZone),
        ];

        for (rcode, parsed_rcode) in valid_rcodes {
            dns_query[3] = rcode << 0;
            let parse_result = DNSHeader::from_udp_datagram(&dns_query);
            assert_eq!(
                parse_result?.response_code, parsed_rcode,
                "query: {:02x?}, rcode: {}, parsed_rcode: {:?}",
                dns_query, rcode, parsed_rcode
            );
        }
        Ok(())
    }
}