/*-
* cdns-rs - a simple sync/async DNS query library
* Copyright (C) 2020  Aleksandr Morozov, RELKOM s.r.o
* Copyright (C) 2021  Aleksandr Morozov
* 
* This program is free software; you can redistribute it and/or
* modify it under the terms of the GNU Lesser General Public
* License as published by the Free Software Foundation; either
* version 3 of the License, or (at your option) any later version.
*
* This program is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
* Lesser General Public License for more details.
* 
* You should have received a copy of the GNU Lesser General Public License
* along with this program; if not, write to the Free Software Foundation,
* Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
*/

use std::collections::LinkedList;
use std::fmt;
use std::net::IpAddr;
use std::time::Duration;

use crate::error::*;
use crate::{write_error, internal_error, internal_error_map};

use super::query::{QDns, QDnsName, QuerySetup, QueryMode};
use super::common::{QType, DnsRdata};


/// Resolves the FQDN to a list of an IPv6 addresses.
/// 
/// # Arguments
/// 
/// * `fqdn` - a domain which is needed to be resolved
/// 
/// # Returns
/// 
/// * [CDnsResult] with vector of [IpAddr]
/// 
/// * [CDnsResult] with error
pub async
fn resolve_aaaa<C>(fqdn: C) -> CDnsResult<Vec<IpAddr>>
where C: AsRef<str>
{
    let res = 
        QDns::query_optional(
            QDnsName::from(fqdn.as_ref()), 
            QType::AAAA, 
            QuerySetup::new(
                false, 
                false, 
                QueryMode::ForceNxQueryAll, 
                Some(Duration::from_secs(3))
            )
        ).await?;

    let dns_qs = 
        match res.into_inner()
        {
            Some(r) => r,
            None => return Ok(Vec::new()),
        };

    let mut ret_list: Vec<IpAddr> = Vec::with_capacity(dns_qs.len());

    for dq in dns_qs
    {
        let ok_res = dq.collect_ok_responses();

        for okr in ok_res
        {
            match okr
            {
                DnsRdata::AAAA{ ip } =>
                {
                    ret_list.push(IpAddr::from(ip));
                },
                _ => {} // skip rest
            }
        }
    }

    return Ok(ret_list);
}

/// Resolves the FQDN to a list of an IPv4 addresses.
/// 
/// # Arguments
/// 
/// * `fqdn` - a domain which is needed to be resolved
/// 
/// # Returns
/// 
/// * [CDnsResult] with vector of [IpAddr]
/// 
/// * [CDnsResult] with error
pub async
fn resolve_a<C>(fqdn: C) -> CDnsResult<Vec<IpAddr>>
where C: AsRef<str>
{
    let res = 
        QDns::query_optional(
            QDnsName::from(fqdn.as_ref()), 
            QType::A, 
            QuerySetup::new(
                false, 
                false, 
                QueryMode::ForceNxQueryAll, 
                Some(Duration::from_secs(3))
            )
        ).await?;

    let dns_qs = 
        match res.into_inner()
        {
            Some(r) => r,
            None => return Ok(Vec::new()),
        };

    let mut ret_list: Vec<IpAddr> = Vec::with_capacity(dns_qs.len());

    for dq in dns_qs
    {
        let ok_res = dq.collect_ok_responses();

        for okr in ok_res
        {
            match okr
            {
                DnsRdata::A{ ip } =>
                {
                    ret_list.push(IpAddr::from(ip));
                },
                _ => {} // skip rest
            }
        }
    }

    return Ok(ret_list);
}

#[derive(Debug)]
pub struct DnsMxResolved
{
    pub preference: u16,
    pub exchange_fqdn: String, 
    pub ip: IpAddr
}

impl fmt::Display for DnsMxResolved
{
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result 
    {
        write!(f, "{} {} {}", self.preference, self.exchange_fqdn, self.ip)    
    }
}

/// Searches for the MX records of the provided domain.
/// 
/// Does not take into account the preference.
/// 
/// The reply (on request) timeout is 3 seconds.
/// 
/// # Arguments
/// 
/// * `domain` - a domain which is required to check.
/// 
/// * `aaaa` - if set to true then AAAA will also be resolved.
/// 
/// # Returns
/// 
/// * [CDnsResult] with a pair of FQDN, IP address the FQDN points to
/// 
/// * [CDnsResult] with an error message
pub async 
fn resolve_mx<D>(domain: D, aaaa: bool) -> CDnsResult<Option<Vec<DnsMxResolved>>>
where D: AsRef<str>
{
    let mut ret_list: LinkedList<DnsMxResolved> = LinkedList::new();

    let res = 
        QDns::query_optional(
            QDnsName::from(domain.as_ref()), 
            QType::MX, 
            QuerySetup::new(
                false, 
                false, 
                QueryMode::ForceNxQueryAll, 
                Some(Duration::from_secs(3))
            )
        ).await?;

    let dns_qs = 
        match res.into_inner()
        {
            Some(r) => r,
            None => return Ok(None),
        };

    for dq in dns_qs
    {
        let ok_res = dq.collect_ok_responses();

        for okr in ok_res
        {
            match okr
            {
                DnsRdata::MX{ preference, exchange } =>
                {
                    // use preference to insert the record in order
                    // exchange sould be CNAME, but may be an IP address if set up incorectly
                    // check if it really FQDN, not IP

                    // send new request to resolve fqdn
                    match resolve_a(&exchange).await
                    {
                        Ok(ip_list) =>
                        {
                            for ip in ip_list
                            {
                                ret_list.push_back(
                                    DnsMxResolved
                                    {
                                        preference: preference,
                                        exchange_fqdn: exchange.clone(),
                                        ip: ip
                                    }
                                );
                            }
                        },
                        Err(e) => 
                        {
                            write_error!("{}", e);
                            continue;
                        }
                    }

                    // check if AAAA is needed
                    if aaaa == true
                    {
                        match resolve_aaaa(&exchange).await
                        {
                            Ok(ip_list) =>
                            {
                                for ip in ip_list
                                {
                                    ret_list.push_back(
                                        DnsMxResolved
                                        {
                                            preference: preference,
                                            exchange_fqdn: exchange.clone(),
                                            ip: ip
                                        }
                                    );
                                }
                            },
                            Err(e) => 
                            {
                                write_error!("{}", e);
                                continue;
                            }
                        }
                    }

                },
                _ => {} // skip rest
            }
        }
    }

    return Ok(Some(ret_list.into_iter().map(|d| d).collect()));
}


#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
async fn test_mx_resolve()
{
    let fqdn_ip = resolve_mx("protonmail.com", true).await;

    assert_eq!(fqdn_ip.is_err(), false);

    let fqdn_ip = fqdn_ip.unwrap();

    assert_eq!(fqdn_ip.is_some(), true);

    let dnsips = fqdn_ip.unwrap();

    for di in dnsips
    {
        println!("{}", di);
    }
    
}