use hickory_proto::op::Message;

/// Errors that can occur when building a DNS response.
#[derive(Debug, thiserror::Error)]
pub enum ResponseBuildError {
    /// Serialization of the DNS message failed.
    #[error("failed to serialize DNS message: {0}")]
    SerializationFailed(#[from] hickory_proto::ProtoError),
}

/// Serialize a DNS [`Message`] into wire-format bytes.
///
/// Used to encode the upstream response before returning it to the
/// DoH client.
pub fn serialize_response(message: &Message) -> Result<Vec<u8>, ResponseBuildError> {
    let bytes = message.to_vec()?;
    Ok(bytes)
}

/// Build a SERVFAIL response for a given query message.
///
/// Used when the upstream resolver is unreachable or returns an error.
pub fn build_servfail(original: &Message) -> Result<Vec<u8>, ResponseBuildError> {
    let mut response = Message::new();
    response.set_id(original.id());
    response.set_response_code(hickory_proto::op::ResponseCode::ServFail);
    response.set_recursion_desired(original.recursion_desired());
    response.set_recursion_available(true);

    // Copy the original question section into the response
    for query in original.queries() {
        response.add_query(query.clone());
    }

    serialize_response(&response)
}

#[cfg(test)]
mod tests {
    use super::*;
    use hickory_proto::op::{Message, Query};
    use hickory_proto::rr::RecordType;

    #[test]
    fn test_serialize_and_roundtrip() {
        let mut msg = Message::new();
        msg.set_id(1234);
        msg.set_response_code(hickory_proto::op::ResponseCode::NoError);

        let bytes = serialize_response(&msg).unwrap();
        assert!(!bytes.is_empty());

        let parsed = Message::from_vec(&bytes).unwrap();
        assert_eq!(parsed.id(), 1234);
    }

    #[test]
    fn test_build_servfail_preserves_id_and_question() {
        let mut query = Message::new();
        query.set_id(5678);
        query.add_query(Query::query(
            hickory_proto::rr::Name::from_ascii("example.com.").unwrap(),
            RecordType::A,
        ));

        let response_bytes = build_servfail(&query).unwrap();
        let response = Message::from_vec(&response_bytes).unwrap();

        assert_eq!(response.id(), 5678);
        assert_eq!(
            response.response_code(),
            hickory_proto::op::ResponseCode::ServFail
        );
        assert_eq!(response.queries().len(), 1);
    }
}
