#![deny(warnings)]

use anyhow::{bail, Context, Result};
use std::convert::TryFrom;
use std::net::IpAddr;
use std::str;

use bytes::{BufMut, BytesMut};
use packed_struct::prelude::*;

use crate::codec::{domain_name, rdata};
use crate::fbs::dns_enums_generated::ResourceType;
use crate::fbs::dns_message_generated::*;

/// Wire protocol type for OPT records, for use in some match statements.
pub const OPT_RESOURCE_TYPE: u16 = 41;

// Header, from RFC2535 section 6.1 or RFC2929 section 2.
//
//                                 1  1  1  1  1  1
//   0  1  2  3  4  5  6  7  8  9  0  1  2  3  4  5
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |                      ID                       |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |QR|   Opcode  |AA|TC|RD|RA| Z|AD|CD|   RCODE   |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |                    QDCOUNT                    |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |                    ANCOUNT                    |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |                    NSCOUNT                    |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |                    ARCOUNT                    |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
#[derive(PackedStruct)]
#[packed_struct(endian = "msb", bit_numbering = "msb0")]
pub struct HeaderBits {
    #[packed_field(bits = "0:15")]
    pub id: u16,

    #[packed_field(bits = "16:16")]
    pub is_response: bool,
    #[packed_field(bits = "17:20")]
    pub op_code: Integer<u8, packed_bits::Bits4>,
    #[packed_field(bits = "21:21")]
    pub authoritative: bool,
    #[packed_field(bits = "22:22")]
    pub truncated: bool,
    #[packed_field(bits = "23:23")]
    pub recursion_desired: bool,
    #[packed_field(bits = "24:24")]
    pub recursion_available: bool,
    #[packed_field(bits = "25:25")]
    pub reserved_9: bool,
    #[packed_field(bits = "26:26")]
    pub authentic_data: bool,
    #[packed_field(bits = "27:27")]
    pub checking_disabled: bool,
    #[packed_field(bits = "28:31")]
    pub response_code: Integer<u8, packed_bits::Bits4>,

    #[packed_field(bits = "32:47")]
    pub question_count: u16,
    #[packed_field(bits = "48:63")]
    pub answer_count: u16,
    #[packed_field(bits = "64:79")]
    pub authority_count: u16,
    #[packed_field(bits = "80:95")]
    pub additional_count: u16,
}

#[derive(Clone, Debug, PartialEq)]
pub struct RecordCounts {
    pub question: u16,
    pub answer: u16,
    pub authority: u16,
    pub additional: u16,
}

impl RecordCounts {
    pub fn new() -> RecordCounts {
        RecordCounts {
            question: 0,
            answer: 0,
            authority: 0,
            additional: 0,
        }
    }
}

/// Writes the message header to the provided buffer.
/// Allows specifying the transaction id since this is frequently overridden.
pub fn write_header_id(message: &Message, id_override: u16, buf: &mut BytesMut) -> Result<()> {
    let header = message
        .header()
        .with_context(|| "required header is missing")?;
    // Ensure that the OPT record, if any, is included in additional_count
    let opt_len = match message.opt() {
        Some(_opt) => 1,
        None => 0,
    };
    write_header_bits(
        HeaderBits {
            id: id_override,

            is_response: header.is_response(),
            op_code: Integer::from(header.op_code()),
            authoritative: header.authoritative(),
            truncated: header.truncated(),
            recursion_desired: header.recursion_desired(),
            recursion_available: header.recursion_available(),
            reserved_9: header.reserved_9(),
            authentic_data: header.authentic_data(),
            checking_disabled: header.checking_disabled(),
            response_code: Integer::from(header.response_code()),

            question_count: u16::try_from(match message.question() {
                Some(vec) => vec.len(),
                None => 0,
            })
            .with_context(|| "question length doesn't fit")?,
            answer_count: u16::try_from(match message.answer() {
                Some(vec) => vec.len(),
                None => 0,
            })
            .with_context(|| "answer length doesn't fit")?,
            authority_count: u16::try_from(match message.authority() {
                Some(vec) => vec.len(),
                None => 0,
            })
            .with_context(|| "authority length doesn't fit")?,
            additional_count: u16::try_from(match message.additional() {
                Some(vec) => vec.len() + opt_len,
                None => opt_len,
            })
            .with_context(|| "additional length doesn't fit")?,
        },
        buf,
    )
}

pub fn write_header_bits(bits: HeaderBits, buf: &mut BytesMut) -> Result<()> {
    let bits_packed = bits.pack();
    buf.reserve(bits_packed.len());
    buf.put_slice(&bits_packed);
    Ok(())
}

/// Returns the header content, along with whether it declares that the rest of the Message is truncated.
/// When this is encountered the message should be abandoned and a different connection method should be used.
/// (For example if the original request was UDP then try TCP)
/// Alternatively returns Ok(None) when there isn't enough data available to parse the header yet.
pub fn read_header<'a>(
    fb_builder: &mut flatbuffers::FlatBufferBuilder<'a>,
    buf: &[u8],
    offset: &mut usize,
) -> Result<Option<(flatbuffers::WIPOffset<Header<'a>>, RecordCounts, bool)>> {
    let bits_size = HeaderBits::packed_bytes();
    if buf.len() < *offset + bits_size {
        // Not enough bytes for the DNS header, wait for more bytes.
        return Ok(None);
    }
    let bits = HeaderBits::unpack_from_slice(&buf[*offset..*offset + bits_size])
        .with_context(|| "couldn't unpack header bits")?;

    let header = Header::create(
        fb_builder,
        &HeaderArgs {
            id: bits.id,
            is_response: bits.is_response,
            op_code: *bits.op_code,
            authoritative: bits.authoritative,
            truncated: bits.truncated,
            recursion_desired: bits.recursion_desired,
            recursion_available: bits.recursion_available,
            reserved_9: bits.reserved_9,
            authentic_data: bits.authentic_data,
            checking_disabled: bits.checking_disabled,
            response_code: *bits.response_code,
        },
    );

    let record_counts = RecordCounts {
        question: bits.question_count,
        answer: bits.answer_count,
        authority: bits.authority_count,
        additional: bits.additional_count,
    };

    *offset += bits_size;
    Ok(Some((header, record_counts, bits.truncated)))
}

#[derive(PackedStruct)]
#[packed_struct(endian = "msb", bit_numbering = "msb0")]
pub struct HeaderIDBits {
    #[packed_field(bits = "0:15")]
    id: u16,
}

/// Directly updates the request ID in a serialized Message.
/// This is a shortcut to avoid needing to reserialize a new message in the event of a retry.
pub fn update_message_id(id: u16, buf: &mut BytesMut, offset: usize) -> Result<()> {
    // The ID is (conveniently) located in the first two bytes of the message.
    let bits_size = HeaderIDBits::packed_bytes();
    let buf_len = buf.len();
    let to_update = buf.get_mut(offset..offset + bits_size).with_context(|| {
        format!(
            "Provided message to update is {} bytes, must be at least {}",
            buf_len, bits_size
        )
    })?;
    let bits = HeaderIDBits { id }.pack();
    for i in 0..bits_size {
        to_update[i] = bits[i];
    }
    Ok(())
}

/// Finds the minimum TTL across all resources in a Message, or None if there are no resources.
/// This is a utility for caching where we want to detect how much the response TTLs should be reduced.
/// For example, if the cache says that the remaining TTL is 20s, but the minimum resource TTL value is 30s,
/// then 10s has passed and all resources should be have TTLs reduced by 10s.
pub fn get_min_ttl(message: &Message) -> Option<u32> {
    let mut min_ttl: Option<u32> = None;

    // Skip header and question(s): No TTLs

    if let Some(resources) = message.answer() {
        for i in 0..resources.len() {
            let ttl = resources.get(i).ttl();
            match min_ttl {
                Some(min_val) => {
                    if ttl < min_val {
                        min_ttl = Some(ttl);
                    }
                }
                None => {
                    min_ttl = Some(ttl);
                }
            }
        }
    }

    if let Some(resources) = message.authority() {
        for i in 0..resources.len() {
            let ttl = resources.get(i).ttl();
            match min_ttl {
                Some(min_val) => {
                    if ttl < min_val {
                        min_ttl = Some(ttl);
                    }
                }
                None => {
                    min_ttl = Some(ttl);
                }
            }
        }
    }

    // Note: any OPT resource is in a separate message.opt() field. OPT resources don't have TTLs.
    if let Some(resources) = message.additional() {
        for i in 0..resources.len() {
            let ttl = resources.get(i).ttl();
            match min_ttl {
                Some(min_val) => {
                    if ttl < min_val {
                        min_ttl = Some(ttl);
                    }
                }
                None => {
                    min_ttl = Some(ttl);
                }
            }
        }
    }

    return min_ttl;
}

// Question, from RFC1035 section 4.1.2
//
//                                 1  1  1  1  1  1
//   0  1  2  3  4  5  6  7  8  9  0  1  2  3  4  5
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |                                               |
// /                     QNAME                     /
// /                                               /
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |                     QTYPE                     |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |                     QCLASS                    |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
#[derive(PackedStruct)]
#[packed_struct(endian = "msb", bit_numbering = "msb0")]
pub struct QuestionBits {
    // name: variable length string
    #[packed_field(bits = "0:15")]
    resource_type: u16,
    #[packed_field(bits = "16:31")]
    resource_class: u16,
}

pub fn write_question(
    question: &Question,
    buf: &mut BytesMut,
    ptr_offsets: &mut domain_name::LabelOffsets,
) -> Result<()> {
    domain_name::write(
        &question.name().with_context(|| "missing question.name")?,
        buf,
        ptr_offsets,
        "question.name",
    )?;

    let bits = QuestionBits {
        resource_type: question.resource_type(),
        resource_class: question.resource_class(),
    }
    .pack();
    buf.reserve(bits.len());
    buf.put_slice(&bits);
    Ok(())
}

pub fn read_question<'a>(
    fb_builder: &mut flatbuffers::FlatBufferBuilder<'a>,
    buf: &[u8],
    offset: &mut usize,
) -> Result<Option<flatbuffers::WIPOffset<Question<'a>>>> {
    let (name_bytes_consumed, name_str) = domain_name::read(buf, *offset, "question.name")?;

    let bits_size = QuestionBits::packed_bytes();
    if buf.len() < *offset + name_bytes_consumed + bits_size {
        // Not enough bytes for the question header, wait for more bytes.
        return Ok(None);
    }
    let bits = QuestionBits::unpack_from_slice(
        &buf[*offset + name_bytes_consumed..*offset + name_bytes_consumed + bits_size],
    )
    .with_context(|| "couldn't unpack question bits")?;

    *offset += name_bytes_consumed + bits_size;

    let args = QuestionArgs {
        name: Some(fb_builder.create_string(name_str.as_str())),
        resource_type: bits.resource_type,
        resource_class: bits.resource_class,
    };
    Ok(Some(Question::create(fb_builder, &args)))
}

// Resource, from RFC1035 section 4.1.3
//
//                                 1  1  1  1  1  1
//   0  1  2  3  4  5  6  7  8  9  0  1  2  3  4  5
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |                                               |
// /                      NAME                     /
// /                                               /
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |                      TYPE                     |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |                     CLASS                     |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |                      TTL                      |
// |                                               |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |                   RDLENGTH                    |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--|
// /                     RDATA                     /
// /                                               /
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+

// Just the initial type field, for reading out and deciding if we're OPT or not.
#[derive(PackedStruct)]
#[packed_struct(endian = "msb", bit_numbering = "msb0")]
pub struct ResourceTypeBits {
    // name: variable length string
    #[packed_field(bits = "0:15")]
    resource_type: u16,
}

// The regular Class/TTL fields used by non-OPT records
#[derive(PackedStruct)]
#[packed_struct(endian = "msb", bit_numbering = "msb0")]
pub struct ResourceClassTTLBits {
    #[packed_field(bits = "0:15")]
    class: u16,
    #[packed_field(bits = "16:47")]
    ttl: u32,
}

// OPT-specific Class/TTL packing, from RFC6891 section 6.1.3:
//                                 1  1  1  1  1  1
//   0  1  2  3  4  5  6  7  8  9  0  1  2  3  4  5
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |                   UDP-SIZE                    |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |     EXTENDED-RCODE    |        VERSION        |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |DO|                    Z                       |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// Mirrors ResourceClassTTLBits with OPT packing
// Members are public to allow access by rdata.rs
#[derive(Default, PackedStruct)]
#[packed_struct(endian = "msb", bit_numbering = "msb0")]
pub struct ResourceClassTTLOPTBits {
    // Class field
    #[packed_field(bits = "0:15")]
    pub udp_size: u16,
    // TTL field
    #[packed_field(bits = "16:23")]
    pub response_code: u8,
    #[packed_field(bits = "24:31")]
    pub version: u8,
    #[packed_field(bits = "32:32")]
    pub dnssec_ok: bool,
    #[packed_field(bits = "33:47")]
    _reserved: ReservedZero<packed_bits::Bits15>,
}

// The RDLength field, initialized to zero then updated after writing the RData
#[derive(PackedStruct)]
#[packed_struct(endian = "msb", bit_numbering = "msb0")]
pub struct ResourceRdataLengthBits {
    #[packed_field(bits = "0:15")]
    rdata_len: u16,
}

pub enum RDataFields<'a> {
    RDATA(Option<ResourceData<'a>>),
    IP(IpAddr),
    SOA(rdata::SOAFields<'a>),
}

/// Values to be used when writing a Resource, instead of using what's in the fbs version.
/// Used when we're constructing our own custom response to something.
pub struct ResourceFields<'a> {
    pub name: &'a str,
    pub resource_type: u16,
    pub resource_class: u16,
    pub ttl: u32,
    pub rdata: RDataFields<'a>,
}

/// Encodes a DNS resource to the provided buf
pub fn write_resource(
    resource: &Resource,
    buf: &mut BytesMut,
    ptr_offsets: &mut domain_name::LabelOffsets,
) -> Result<()> {
    write_resource_fields(
        &ResourceFields {
            name: resource.name().with_context(|| "missing resource.name")?,
            resource_type: resource.resource_type(),
            resource_class: resource.resource_class(),
            ttl: resource.ttl(),
            rdata: RDataFields::RDATA(resource.rdata()),
        },
        buf,
        ptr_offsets,
    )
}

/// Encodes a resource to the provided buf, with the TTL reduced by the specified amount
pub fn write_resource_ttl_subtract(
    resource: &Resource,
    ttl_subtract: u32,
    buf: &mut BytesMut,
    ptr_offsets: &mut domain_name::LabelOffsets,
) -> Result<()> {
    write_resource_fields(
        &ResourceFields {
            name: resource.name().with_context(|| "missing resource.name")?,
            resource_type: resource.resource_type(),
            resource_class: resource.resource_class(),
            ttl: resource.ttl() - ttl_subtract,
            rdata: RDataFields::RDATA(resource.rdata()),
        },
        buf,
        ptr_offsets,
    )
}

/// Encodes a resource fields to the provided buf, with support for overriding field values
pub fn write_resource_fields(
    resource: &ResourceFields,
    buf: &mut BytesMut,
    ptr_offsets: &mut domain_name::LabelOffsets,
) -> Result<()> {
    domain_name::write(resource.name, buf, ptr_offsets, "resource.name")?;

    let type_bits = ResourceTypeBits {
        resource_type: resource.resource_type,
    }
    .pack();

    let classttl_bits = ResourceClassTTLBits {
        class: resource.resource_class,
        ttl: resource.ttl,
    }
    .pack();

    let rdlen_zero_bits = ResourceRdataLengthBits {
        rdata_len: 0 as u16, // To be updated below
    }
    .pack();

    buf.reserve(type_bits.len() + classttl_bits.len() + rdlen_zero_bits.len());
    buf.put_slice(&type_bits);
    buf.put_slice(&classttl_bits);
    let rdata_len_offset = buf.len();
    buf.put_slice(&rdlen_zero_bits);
    let rdata_offset = buf.len();

    // We write directly to 'buf' to ensure that any domain-name offsets are relative to the start
    // of the full 'buf'. Afterwards we go back and update the rdata_len to reflect the size of
    // what was written.
    match &resource.rdata {
        RDataFields::RDATA(None) => {}
        RDataFields::RDATA(Some(rdata)) => {
            rdata::write_rdata(resource.resource_type, &rdata, buf, ptr_offsets)?
        }
        RDataFields::IP(IpAddr::V4(ip)) => rdata::write_a_ip(&ip, buf)?,
        RDataFields::IP(IpAddr::V6(ip)) => rdata::write_aaaa_ip(&ip, buf)?,
        RDataFields::SOA(soa) => rdata::write_soa_fields(soa, buf, ptr_offsets)?,
    }

    if buf.len() > rdata_offset {
        // We wrote some rdata, so go back and update the rdata length
        let rdlen_updated_bits = ResourceRdataLengthBits {
            rdata_len: u16::try_from(buf.len() - rdata_offset)
                .with_context(|| "rdata.length doesn't fit")?,
        }
        .pack();
        let buf_len = buf.len();
        let bits = buf
            .get_mut(rdata_len_offset..rdata_offset)
            .with_context(|| {
                format!(
                    "failed to get bytes {}..{} from buffer with length {}",
                    rdata_len_offset, rdata_offset, buf_len
                )
            })?;
        // Overwrite the previously appended rdlen_zero_bits with rdlen_updated_bits
        for i in 0..(rdata_offset - rdata_len_offset) {
            bits[i] = rdlen_updated_bits[i];
        }
    }

    Ok(())
}

/// Reads a resource in the common case when OPT records are not expected.
pub fn read_resource_non_opt<'a>(
    fb_builder: &mut flatbuffers::FlatBufferBuilder<'a>,
    buf: &[u8],
    offset: &mut usize,
) -> Result<Option<flatbuffers::WIPOffset<Resource<'a>>>> {
    match read_resource_name_type(buf, offset).context("Failed to read non-OPT resource")? {
        Some((_name, OPT_RESOURCE_TYPE)) => bail!("Got OPT resource in unexpected part of message"),
        Some((name, resource_type)) => {
            match read_resource_remainder(fb_builder, buf, offset, name, resource_type)? {
                Some(resource) => Ok(Some(resource)),
                None => Ok(None),
            }
        }
        None => Ok(None),
    }
}

/// Reads the resource name and type, or None if not enough data is available.
pub fn read_resource_name_type<'a>(
    buf: &[u8],
    offset: &mut usize,
) -> Result<Option<(String, u16)>> {
    let (mut total_bytes_consumed, name_str) = domain_name::read(buf, *offset, "resource.name")?;

    // Get the type field first: Determine whether this is an OPT record
    let type_bits_size = ResourceTypeBits::packed_bytes();
    if buf.len() < *offset + total_bytes_consumed + type_bits_size {
        // Not enough bytes for the resource type field, wait for more bytes.
        return Ok(None);
    }
    let resource_type = ResourceTypeBits::unpack_from_slice(
        &buf[*offset + total_bytes_consumed..*offset + total_bytes_consumed + type_bits_size],
    )
    .with_context(|| "couldn't unpack resource prelude bits")?
    .resource_type;
    total_bytes_consumed += type_bits_size;

    *offset += total_bytes_consumed;

    Ok(Some((name_str, resource_type)))
}

/// Reads the rest of an OPT resource.
/// This should be called after read_resource_name_type(), which provides the type of resource being read.
pub fn read_resource_remainder_opt<'a>(
    fb_builder: &mut flatbuffers::FlatBufferBuilder<'a>,
    buf: &[u8],
    offset: &mut usize,
) -> Result<Option<flatbuffers::WIPOffset<OPT<'a>>>> {
    // OPT/EDNS(0) hax: Custom packing with OPT-specific fields
    let opt_class_ttl_bits_size = ResourceClassTTLOPTBits::packed_bytes();
    if buf.len() < *offset + opt_class_ttl_bits_size {
        // Not enough bytes for the OPT class/ttl fields, wait for more bytes.
        return Ok(None);
    }
    let opt_class_ttl_bits = ResourceClassTTLOPTBits::unpack_from_slice(
        &buf[*offset..*offset + opt_class_ttl_bits_size],
    )
    .with_context(|| "couldn't unpack OPT resource custom class/ttl bits")?;
    let mut total_bytes_consumed = opt_class_ttl_bits_size;

    let rdata_location = read_rdata_location(
        buf,
        *offset + total_bytes_consumed,
        &mut total_bytes_consumed,
    )
    .context("Failed to read OPT rdata location")?;
    match rdata_location {
        Some((rdata_offset, rdata_len)) => {
            // Pass the OPT-specific data replacing the class/TTL bits, to be stored in the OPT object
            let opt =
                rdata::read_opt(fb_builder, buf, rdata_offset, rdata_len, opt_class_ttl_bits)?;
            total_bytes_consumed += rdata_len;
            *offset += total_bytes_consumed;
            Ok(Some(opt))
        }
        None => {
            // Not enough bytes for the size info (or the expected rdata), wait for more bytes.
            Ok(None)
        }
    }
}

/// Reads the rest of a non-OPT resource.
/// This should be called after read_resource_name_type(), which provides the type of resource being read.
pub fn read_resource_remainder<'a>(
    fb_builder: &mut flatbuffers::FlatBufferBuilder<'a>,
    buf: &[u8],
    offset: &mut usize,
    name_str: String,
    resource_type: u16,
) -> Result<Option<flatbuffers::WIPOffset<Resource<'a>>>> {
    // Non-OPT resource: regular class and TTL fields
    let class_ttl_bits_size = ResourceClassTTLBits::packed_bytes();
    if buf.len() < *offset + class_ttl_bits_size {
        // Not enough bytes for the class/ttl fields, wait for more bytes.
        return Ok(None);
    }
    let class_ttl_bits =
        ResourceClassTTLBits::unpack_from_slice(&buf[*offset..*offset + class_ttl_bits_size])
            .with_context(|| "couldn't unpack resource class/ttl bits")?;
    let mut total_bytes_consumed = class_ttl_bits_size;

    let rdata_location = read_rdata_location(
        buf,
        *offset + total_bytes_consumed,
        &mut total_bytes_consumed,
    )
    .with_context(|| {
        format!(
            "Failed to read {:?} resource rdata location for {}",
            resource_type, name_str
        )
    })?;
    match rdata_location {
        Some((rdata_offset, rdata_len)) => {
            let rdata = rdata::read_rdata(fb_builder, buf, resource_type, rdata_offset, rdata_len)
                .with_context(|| {
                    format!(
                        "Failed to read {:?} resource for {}",
                        resource_type, name_str
                    )
                })?;
            total_bytes_consumed += rdata_len;
            *offset += total_bytes_consumed;
            let args = ResourceArgs {
                name: Some(fb_builder.create_string(name_str.as_str())),
                resource_type,
                resource_class: class_ttl_bits.class,
                ttl: class_ttl_bits.ttl,
                rdata: Some(ResourceData::create(fb_builder, &rdata)),
            };
            Ok(Some(Resource::create(fb_builder, &args)))
        }
        None => {
            // Not enough bytes for the size info (or the expected rdata), wait for more bytes.
            Ok(None)
        }
    }
}

/// Returns the offset and length of the rdata available in the buffer,
/// or None if the buffer is incomplete
fn read_rdata_location<'a>(
    buf: &[u8],
    offset: usize,
    total_bytes_consumed: &mut usize,
) -> Result<Option<(usize, usize)>> {
    let rdata_length_bits_size = ResourceRdataLengthBits::packed_bytes();
    let rdata_offset = offset + rdata_length_bits_size;
    if buf.len() < rdata_offset {
        // Not enough bytes for the rdata length field, wait for more bytes.
        return Ok(None);
    }
    let rdata_length_bits =
        ResourceRdataLengthBits::unpack_from_slice(&buf[offset..offset + rdata_length_bits_size])
            .with_context(|| "couldn't unpack rdata length bits")?;

    let rdata_len = rdata_length_bits.rdata_len as usize;
    if buf.len() < rdata_offset + rdata_len {
        // Not enough bytes for the rdata content, wait for more bytes.
        return Ok(None);
    }

    *total_bytes_consumed += rdata_length_bits_size;
    Ok(Some((rdata_offset, rdata_len)))
}

pub fn write_opt(opt: &OPT, udp_size_override: Option<u16>, buf: &mut BytesMut) -> Result<()> {
    domain_name::write_nopointer(".", buf, "opt.name")?;

    let type_bits = ResourceTypeBits {
        resource_type: ResourceType::TYPE_OPT as u16,
    }
    .pack();

    let classttl_bits = ResourceClassTTLOPTBits {
        udp_size: match udp_size_override {
            Some(udp_size) => udp_size,
            None => opt.udp_size(),
        },
        response_code: opt.response_code(),
        version: opt.response_code(),
        dnssec_ok: opt.dnssec_ok(),
        ..ResourceClassTTLOPTBits::default() // for setting the _reserved field
    }
    .pack();

    let rdlen_zero_bits = ResourceRdataLengthBits {
        rdata_len: 0 as u16, // To be updated below
    }
    .pack();

    buf.reserve(type_bits.len() + classttl_bits.len() + rdlen_zero_bits.len());
    buf.put_slice(&type_bits);
    buf.put_slice(&classttl_bits);
    let rdata_len_offset = buf.len();
    buf.put_slice(&rdlen_zero_bits);
    let rdata_offset = buf.len();

    rdata::write_opt(opt, buf)?;

    if buf.len() > rdata_offset {
        // We wrote some rdata, so go back and update the rdata length
        let rdlen_updated_bits = ResourceRdataLengthBits {
            rdata_len: u16::try_from(buf.len() - rdata_offset)
                .with_context(|| "rdata.length doesn't fit")?,
        }
        .pack();
        let buf_len = buf.len();
        let bits = buf
            .get_mut(rdata_len_offset..rdata_offset)
            .with_context(|| {
                format!(
                    "failed to get bytes {}..{} from buffer with length {}",
                    rdata_len_offset, rdata_offset, buf_len
                )
            })?;
        // Overwrite the previously appended rdlen_zero_bits with rdlen_updated_bits
        for i in 0..(rdata_offset - rdata_len_offset) {
            bits[i] = rdlen_updated_bits[i];
        }
    }
    Ok(())
}