rsubdomain 1.2.11

use std::collections::HashMap;
use std::net::IpAddr;
use std::sync::Arc;

use trust_dns_resolver::config::*;
use trust_dns_resolver::proto::rr::{RData, RecordType};
use trust_dns_resolver::TokioAsyncResolver;

/// DNS记录类型
#[derive(Debug, Clone)]
pub enum DnsRecord {
    A(String),       // IPv4地址
    AAAA(String),    // IPv6地址
    CNAME(String),   // 别名
    NS(String),      // 名称服务器
    MX(u16, String), // 邮件交换器 (优先级, 主机名)
    TXT(String),     // 文本记录
    SOA(String),     // 授权开始
    PTR(String),     // 指针记录
}

/// DNS解析结果
#[derive(Debug, Clone)]
pub struct DnsResolveResult {
    pub domain: String,
    pub records: HashMap<String, Vec<DnsRecord>>,
    pub has_records: bool,
}

/// DNS解析器
pub struct DnsResolver {
    resolver: TokioAsyncResolver,
}

impl DnsResolver {
    pub async fn new() -> Result<Self, Box<dyn std::error::Error>> {
        Self::new_with_resolvers(&[]).await
    }

    pub async fn new_with_resolvers(
        resolvers: &[String],
    ) -> Result<Self, Box<dyn std::error::Error>> {
        let resolver =
            TokioAsyncResolver::tokio(build_resolver_config(resolvers)?, ResolverOpts::default());
        Ok(DnsResolver { resolver })
    }

    /// 解析域名的所有记录类型
    pub async fn resolve_all_records(&self, domain: &str) -> DnsResolveResult {
        resolve_all_records_with(&self.resolver, domain).await
    }

    /// 批量解析域名
    pub async fn resolve_domains(&self, domains: Vec<String>) -> Vec<DnsResolveResult> {
        let mut results = Vec::new();
        let semaphore = Arc::new(tokio::sync::Semaphore::new(20));
        let mut tasks = Vec::new();

        for domain in domains {
            let permit = Arc::clone(&semaphore);
            let resolver = self.resolver.clone();

            let task = tokio::spawn(async move {
                match permit.acquire().await {
                    Ok(_permit) => resolve_all_records_with(&resolver, &domain).await,
                    Err(e) => {
                        eprintln!("错误: 获取信号量失败: {}", e);
                        DnsResolveResult {
                            domain,
                            records: HashMap::new(),
                            has_records: false,
                        }
                    }
                }
            });
            tasks.push(task);
        }

        for task in tasks {
            if let Ok(result) = task.await {
                results.push(result);
            }
        }

        results
    }

    /// 显示解析结果
    pub fn display_results(&self, results: &[DnsResolveResult]) {
        println!("=== DNS解析结果 ===");
        for result in results {
            if result.has_records {
                println!("域名: {}", result.domain);
                for (record_type, records) in &result.records {
                    println!("  {}:", record_type);
                    for record in records {
                        match record {
                            DnsRecord::A(ip) => println!("    {}", ip),
                            DnsRecord::AAAA(ip) => println!("    {}", ip),
                            DnsRecord::CNAME(cname) => println!("    {}", cname),
                            DnsRecord::NS(ns) => println!("    {}", ns),
                            DnsRecord::MX(priority, host) => println!("    {} {}", priority, host),
                            DnsRecord::TXT(txt) => println!("    \"{}\"", txt),
                            DnsRecord::SOA(soa) => println!("    {}", soa),
                            DnsRecord::PTR(ptr) => println!("    {}", ptr),
                        }
                    }
                }
                println!();
            }
        }
    }

    /// 简单的A记录解析
    pub async fn resolve_a_record(&self, domain: &str) -> Option<String> {
        if let Ok(response) = self.resolver.lookup_ip(domain).await {
            for ip in response.iter() {
                if let IpAddr::V4(ipv4) = ip {
                    return Some(ipv4.to_string());
                }
            }
        }
        None
    }
}

fn build_resolver_config(
    resolvers: &[String],
) -> Result<ResolverConfig, Box<dyn std::error::Error>> {
    let ips: Vec<IpAddr> = resolvers
        .iter()
        .map(|resolver| resolver.parse())
        .collect::<Result<Vec<IpAddr>, _>>()?;

    if ips.is_empty() {
        return Ok(ResolverConfig::default());
    }

    Ok(ResolverConfig::from_parts(
        None,
        vec![],
        NameServerConfigGroup::from_ips_clear(&ips, 53, true),
    ))
}

async fn resolve_all_records_with(resolver: &TokioAsyncResolver, domain: &str) -> DnsResolveResult {
    let mut records = HashMap::new();
    let mut has_records = false;

    if let Ok(response) = resolver.lookup_ip(domain).await {
        let mut a_records = Vec::new();
        for ip in response.iter() {
            match ip {
                IpAddr::V4(ipv4) => {
                    a_records.push(DnsRecord::A(ipv4.to_string()));
                    has_records = true;
                }
                IpAddr::V6(ipv6) => {
                    a_records.push(DnsRecord::AAAA(ipv6.to_string()));
                    has_records = true;
                }
            }
        }
        if !a_records.is_empty() {
            records.insert("A/AAAA".to_string(), a_records);
        }
    }

    if let Ok(response) = resolver.lookup(domain, RecordType::CNAME).await {
        let mut cname_records = Vec::new();
        for record in response.iter() {
            if let RData::CNAME(cname) = record {
                cname_records.push(DnsRecord::CNAME(cname.to_string()));
                has_records = true;
            }
        }
        if !cname_records.is_empty() {
            records.insert("CNAME".to_string(), cname_records);
        }
    }

    if let Ok(response) = resolver.lookup(domain, RecordType::NS).await {
        let mut ns_records = Vec::new();
        for record in response.iter() {
            if let RData::NS(ns) = record {
                ns_records.push(DnsRecord::NS(ns.to_string()));
                has_records = true;
            }
        }
        if !ns_records.is_empty() {
            records.insert("NS".to_string(), ns_records);
        }
    }

    if let Ok(response) = resolver.lookup(domain, RecordType::MX).await {
        let mut mx_records = Vec::new();
        for record in response.iter() {
            if let RData::MX(mx) = record {
                mx_records.push(DnsRecord::MX(mx.preference(), mx.exchange().to_string()));
                has_records = true;
            }
        }
        if !mx_records.is_empty() {
            records.insert("MX".to_string(), mx_records);
        }
    }

    if let Ok(response) = resolver.lookup(domain, RecordType::TXT).await {
        let mut txt_records = Vec::new();
        for record in response.iter() {
            if let RData::TXT(txt) = record {
                let txt_data = txt
                    .iter()
                    .map(|bytes| String::from_utf8_lossy(bytes).to_string())
                    .collect::<Vec<_>>()
                    .join("");
                txt_records.push(DnsRecord::TXT(txt_data));
                has_records = true;
            }
        }
        if !txt_records.is_empty() {
            records.insert("TXT".to_string(), txt_records);
        }
    }

    if let Ok(response) = resolver.lookup(domain, RecordType::SOA).await {
        let mut soa_records = Vec::new();
        for record in response.iter() {
            if let RData::SOA(soa) = record {
                soa_records.push(DnsRecord::SOA(format!("{} {}", soa.mname(), soa.rname())));
                has_records = true;
            }
        }
        if !soa_records.is_empty() {
            records.insert("SOA".to_string(), soa_records);
        }
    }

    DnsResolveResult {
        domain: domain.to_string(),
        records,
        has_records,
    }
}