/*

 * Copyright (C) 2015 Benjamin Fry <benjaminfry@me.com>

 *

 * Licensed under the Apache License, Version 2.0 (the "License");

 * you may not use this file except in compliance with the License.

 * You may obtain a copy of the License at

 *

 *     https://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

//! Query struct for looking up resource records

#[cfg(test)]

use alloc::vec::Vec;

use core::fmt::{self, Display, Formatter};

#[cfg(feature = "serde")]

use serde::{Deserialize, Serialize};

use crate::error::*;

use crate::rr::dns_class::DNSClass;

use crate::rr::domain::Name;

use crate::rr::record_type::RecordType;

use crate::rr::{Record, RecordData};

use crate::serialize::binary::*;

#[cfg(feature = "mdns")]

/// From [RFC 6762](https://tools.ietf.org/html/rfc6762#section-5.4)

/// ```text

// To avoid large floods of potentially unnecessary responses in these

// cases, Multicast DNS defines the top bit in the class field of a DNS

// question as the unicast-response bit.

/// ```

const MDNS_UNICAST_RESPONSE: u16 = 1 << 15;

/// Query struct for looking up resource records, basically a resource record without RDATA.

///

/// [RFC 1035, DOMAIN NAMES - IMPLEMENTATION AND SPECIFICATION, November 1987](https://tools.ietf.org/html/rfc1035)

///

/// ```text

/// 4.1.2. Question section format

///

/// The question section is used to carry the "question" in most queries,

/// i.e., the parameters that define what is being asked.  The section

/// contains QDCOUNT (usually 1) entries, each of the following format:

///

///                                     1  1  1  1  1  1

///       0  1  2  3  4  5  6  7  8  9  0  1  2  3  4  5

///     +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+

///     |                                               |

///     /                     QNAME / ZNAME             /

///     /                                               /

///     +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+

///     |                     QTYPE / ZTYPE             |

///     +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+

///     |                     QCLASS / ZCLASS           |

///     +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+

///

/// ```

#[derive(Clone, Debug, Eq, Hash, PartialEq)]

#[cfg_attr(feature = "serde", derive(Deserialize, Serialize))]

#[non_exhaustive]

pub struct Query {

    /// QNAME

    pub name: Name,

    /// QTYPE

    pub query_type: RecordType,

    /// QCLASS

    pub query_class: DNSClass,

    /// mDNS unicast-response bit set or not

    #[cfg(feature = "mdns")]

    pub mdns_unicast_response: bool,

}

impl Query {

    /// Return a default query with an empty name and A, IN for the query_type and query_class

    pub fn root() -> Self {

        Self::new(Name::root(), RecordType::A)

    }

    /// Create a new query from name and type, class defaults to IN

    pub fn new(name: Name, query_type: RecordType) -> Self {

        Self {

            name,

            query_type,

            query_class: DNSClass::IN,

            #[cfg(feature = "mdns")]

            mdns_unicast_response: false,

        }

    }

    /// replaces name with the new name

    pub fn set_name(&mut self, name: Name) -> &mut Self {

        self.name = name;

        self

    }

    /// Specify the RecordType being queried

    pub fn set_query_type(&mut self, query_type: RecordType) -> &mut Self {

        self.query_type = query_type;

        self

    }

    /// Specify÷ the DNS class of the Query, almost always IN

    pub fn set_query_class(&mut self, query_class: DNSClass) -> &mut Self {

        self.query_class = query_class;

        self

    }

    /// Changes mDNS unicast-response bit

    /// See [RFC 6762](https://tools.ietf.org/html/rfc6762#section-5.4)

    #[cfg(feature = "mdns")]

    pub fn set_mdns_unicast_response(&mut self, flag: bool) -> &mut Self {

        self.mdns_unicast_response = flag;

        self

    }

    /// Determines if a record is responsive to this

    pub fn matches_record<T: RecordData>(&self, record: &Record<T>) -> bool {

        self.name == record.name

            && (self.query_type == record.record_type()

                || record.record_type() == RecordType::CNAME)

            && self.query_class == record.dns_class

    }

}

impl BinEncodable for Query {

    fn emit(&self, encoder: &mut BinEncoder<'_>) -> ProtoResult<()> {

        self.name.emit(encoder)?;

        self.query_type.emit(encoder)?;

        #[cfg(not(feature = "mdns"))]

        self.query_class.emit(encoder)?;

        #[cfg(feature = "mdns")]

        {

            if self.mdns_unicast_response {

                (u16::from(self.query_class) | MDNS_UNICAST_RESPONSE).emit(encoder)?;

            } else {

                self.query_class.emit(encoder)?;

            }

        }

        Ok(())

    }

}

impl<'r> BinDecodable<'r> for Query {

    fn read(decoder: &mut BinDecoder<'r>) -> Result<Self, DecodeError> {

        let name = Name::read(decoder)?;

        let query_type = RecordType::read(decoder)?;

        #[cfg(feature = "mdns")]

        let mut mdns_unicast_response = false;

        #[cfg(not(feature = "mdns"))]

        let query_class = DNSClass::read(decoder)?;

        #[cfg(feature = "mdns")]

        let query_class = {

            let query_class_value =

                decoder.read_u16()?.unverified(/*DNSClass::from_u16 will verify the value*/);

            if query_class_value & MDNS_UNICAST_RESPONSE > 0 {

                mdns_unicast_response = true;

                DNSClass::from(query_class_value & !MDNS_UNICAST_RESPONSE)

            } else {

                DNSClass::from(query_class_value)

            }

        };

        Ok(Self {

            name,

            query_type,

            query_class,

            #[cfg(feature = "mdns")]

            mdns_unicast_response,

        })

    }

}

impl Display for Query {

    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), fmt::Error> {

        #[cfg(not(feature = "mdns"))]

        {

            write!(

                f,

                "{name} {class} {ty}",

                name = self.name,

                class = self.query_class,

                ty = self.query_type,

            )

        }

        #[cfg(feature = "mdns")]

        {

            write!(

                f,

                "{name} {class} {ty}; mdns_unicast_response: {mdns}",

                name = self.name,

                class = self.query_class,

                ty = self.query_type,

                mdns = self.mdns_unicast_response

            )

        }

    }

}

#[cfg(test)]

mod tests {

    use super::*;

    #[test]

    fn test_read_and_emit() {

        let expect = Query {

            name: Name::from_ascii("WWW.example.com.").unwrap(),

            query_type: RecordType::AAAA,

            query_class: DNSClass::IN,

            ..Query::root()

        };

        let mut byte_vec: Vec<u8> = Vec::with_capacity(512);

        {

            let mut encoder = BinEncoder::new(&mut byte_vec);

            expect.emit(&mut encoder).unwrap();

        }

        let mut decoder = BinDecoder::new(&byte_vec);

        let got = Query::read(&mut decoder).unwrap();

        assert_eq!(got, expect);

    }

    #[cfg(feature = "mdns")]

    #[test]

    fn test_mdns_unicast_response_bit_handling() {

        const QCLASS_OFFSET: usize = 1 /* empty name */ +

        size_of::<u16>() /* query_type */;

        let mut query = Query::root();

        query.set_mdns_unicast_response(true);

        let mut vec_bytes: Vec<u8> = Vec::with_capacity(512);

        {

            let mut encoder = BinEncoder::new(&mut vec_bytes);

            query.emit(&mut encoder).unwrap();

            let query_class_slice = encoder.slice_of(QCLASS_OFFSET, QCLASS_OFFSET + 2);

            assert_eq!(query_class_slice, &[0x80, 0x01]);

        }

        let mut decoder = BinDecoder::new(&vec_bytes);

        let got = Query::read(&mut decoder).unwrap();

        assert_eq!(got.query_class, DNSClass::IN);

        assert!(got.mdns_unicast_response);

    }

}