//! DNS record verification, resolves hostnames and compares against expected values.

use std::collections::HashSet;

use hickory_resolver::proto::rr::RData;
use hickory_resolver::{ResolverBuilder, TokioResolver};
use tracing::instrument;

use crate::config::DnsRecord;
use crate::types::{DnsCheckResult, DnsRecordType};

/// Resolve DNS records and compare against expected values.
/// Returns one `DnsCheckResult` per `DnsRecord` in the input.
#[instrument(skip_all)]
pub async fn check_dns(target: &str, records: &[DnsRecord]) -> Vec<DnsCheckResult> {
    let resolver = match TokioResolver::builder_tokio().and_then(ResolverBuilder::build) {
        Ok(resolver) => resolver,
        Err(e) => {
            return records
                .iter()
                .map(|r| DnsCheckResult {
                    target: target.to_string(),
                    name: r.name.clone(),
                    record_type: r.record_type,
                    expected: r.expected.clone(),
                    actual: vec![],
                    matches: false,
                    checked_at: chrono::Utc::now().to_rfc3339(),
                    error: Some(format!("failed to create resolver: {e}")),
                })
                .collect();
        }
    };

    let mut results = Vec::with_capacity(records.len());
    for record in records {
        let result = resolve_record(target, &resolver, record).await;
        results.push(result);
    }
    results
}

async fn resolve_record(
    target: &str,
    resolver: &TokioResolver,
    record: &DnsRecord,
) -> DnsCheckResult {
    let now = chrono::Utc::now().to_rfc3339();

    let actual = match record.record_type {
        DnsRecordType::A => resolve_a(resolver, &record.name).await,
        DnsRecordType::Aaaa => resolve_aaaa(resolver, &record.name).await,
        DnsRecordType::Cname => resolve_cname(resolver, &record.name).await,
        DnsRecordType::Mx => resolve_mx(resolver, &record.name).await,
        DnsRecordType::Txt => resolve_txt(resolver, &record.name).await,
    };

    match actual {
        Ok(actual_values) => {
            let matches = check_match(&record.expected, &actual_values);
            DnsCheckResult {
                target: target.to_string(),
                name: record.name.clone(),
                record_type: record.record_type,
                expected: record.expected.clone(),
                actual: actual_values,
                matches,
                checked_at: now,
                error: None,
            }
        }
        Err(e) => DnsCheckResult {
            target: target.to_string(),
            name: record.name.clone(),
            record_type: record.record_type,
            expected: record.expected.clone(),
            actual: vec![],
            matches: false,
            checked_at: now,
            error: Some(e),
        },
    }
}

/// Check if all expected values are found in actual (expected ⊆ actual).
pub fn check_match(expected: &[String], actual: &[String]) -> bool {
    let actual_set: HashSet<&str> = actual.iter().map(std::string::String::as_str).collect();
    expected.iter().all(|e| actual_set.contains(e.as_str()))
}

async fn resolve_a(resolver: &TokioResolver, name: &str) -> Result<Vec<String>, String> {
    let response = resolver
        .ipv4_lookup(name)
        .await
        .map_err(|e| format!("A lookup failed for {name}: {e}"))?;
    Ok(rdata(&response)
        .filter_map(|d| match d {
            RData::A(a) => Some(a.0.to_string()),
            _ => None,
        })
        .collect())
}

async fn resolve_aaaa(resolver: &TokioResolver, name: &str) -> Result<Vec<String>, String> {
    let response = resolver
        .ipv6_lookup(name)
        .await
        .map_err(|e| format!("AAAA lookup failed for {name}: {e}"))?;
    Ok(rdata(&response)
        .filter_map(|d| match d {
            RData::AAAA(aaaa) => Some(aaaa.0.to_string()),
            _ => None,
        })
        .collect())
}

async fn resolve_cname(resolver: &TokioResolver, name: &str) -> Result<Vec<String>, String> {
    let response = resolver
        .lookup(name, hickory_resolver::proto::rr::RecordType::CNAME)
        .await
        .map_err(|e| format!("CNAME lookup failed for {name}: {e}"))?;
    Ok(rdata(&response)
        .filter_map(|d| match d {
            RData::CNAME(c) => Some(trim_root(&c.0.to_string())),
            _ => None,
        })
        .collect())
}

async fn resolve_mx(resolver: &TokioResolver, name: &str) -> Result<Vec<String>, String> {
    let response = resolver
        .mx_lookup(name)
        .await
        .map_err(|e| format!("MX lookup failed for {name}: {e}"))?;
    Ok(rdata(&response)
        .filter_map(|d| match d {
            RData::MX(mx) => Some(trim_root(&mx.exchange.to_string())),
            _ => None,
        })
        .collect())
}

async fn resolve_txt(resolver: &TokioResolver, name: &str) -> Result<Vec<String>, String> {
    let response = resolver
        .txt_lookup(name)
        .await
        .map_err(|e| format!("TXT lookup failed for {name}: {e}"))?;
    Ok(rdata(&response)
        .filter_map(|d| match d {
            RData::TXT(txt) => Some(txt.to_string()),
            _ => None,
        })
        .collect())
}

/// Answer-section rdata for a lookup. Typed lookups return a plain `Lookup` as of
/// hickory 0.26, so every caller filters the answers down to the variant it asked for.
fn rdata(response: &hickory_resolver::lookup::Lookup) -> impl Iterator<Item = &RData> {
    response.answers().iter().map(|record| &record.data)
}

fn trim_root(name: &str) -> String {
    name.trim_end_matches('.').to_string()
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn check_match_exact() {
        assert!(check_match(
            &["1.2.3.4".to_string()],
            &["1.2.3.4".to_string()],
        ));
    }

    #[test]
    fn check_match_subset() {
        assert!(check_match(
            &["1.2.3.4".to_string()],
            &["1.2.3.4".to_string(), "5.6.7.8".to_string()],
        ));
    }

    #[test]
    fn check_match_mismatch() {
        assert!(!check_match(
            &["1.2.3.4".to_string()],
            &["5.6.7.8".to_string()],
        ));
    }

    #[test]
    fn check_match_empty_expected() {
        assert!(check_match(&[], &["1.2.3.4".to_string()]));
    }

    #[test]
    fn check_match_empty_actual() {
        assert!(!check_match(&["1.2.3.4".to_string()], &[]));
    }

    #[test]
    fn check_match_order_independent() {
        assert!(check_match(
            &["b".to_string(), "a".to_string()],
            &["a".to_string(), "b".to_string(), "c".to_string()],
        ));
    }

    #[test]
    fn check_match_multiple_expected_all_present() {
        assert!(check_match(
            &["1.2.3.4".to_string(), "5.6.7.8".to_string()],
            &["5.6.7.8".to_string(), "1.2.3.4".to_string()],
        ));
    }

    #[test]
    fn check_match_multiple_expected_one_missing() {
        assert!(!check_match(
            &["1.2.3.4".to_string(), "5.6.7.8".to_string()],
            &["1.2.3.4".to_string()],
        ));
    }
}