/*
 * This module wraps the trust_dns_resolver library
 * to generate results thaat are ready for serializing
 * or templating.
 * It does not aim to be reusable for any other purpose,
 * the trust_dns_resolver library already does that.
 */

use trust_dns_proto::op::response_code::ResponseCode;
use trust_dns_proto::rr::{
	RData,
	record_type::RecordType,
};
use trust_dns_resolver::{
	error::ResolveError,
	error::ResolveErrorKind,
	lookup::Lookup,
	Name,
	TokioAsyncResolver,
};

use tokio::join;

use std::net::IpAddr;


/* Data Structures */

#[derive(serde::Deserialize, serde::Serialize, Default, Clone)]
pub struct DnsLookupResult {
	pub a:     Option<Vec<IpAddr>>,
	pub aaaa:  Option<Vec<IpAddr>>,
	pub aname: Option<Vec<String>>,
	pub cname: Option<Vec<String>>,
	pub mx:    Option<Vec<MxRecord>>,
	pub ns:    Option<Vec<String>>,
	pub soa:   Option<Vec<SoaRecord>>,
	pub txt:   Option<Vec<String>>,
	pub srv:   Option<Vec<SrvRecord>>,
	pub caa:   Option<Vec<String>>,
	pub other_error:     bool,
	pub dns_error:       bool,
	pub nxdomain:        bool,
	pub timeout:         bool,
	pub invalid_name:    bool,
	pub unkown_resolver: bool,
}

#[derive(serde::Deserialize, serde::Serialize, Clone, PartialEq)]
pub struct MxRecord {
	preference: u16,
	exchange:   String,
}

#[derive(serde::Deserialize, serde::Serialize, Clone, PartialEq)]
pub struct SoaRecord {
	mname:   String,
	rname:   String,
	serial:  u32,
	refresh: i32,
	retry:   i32,
	expire:  i32,
	minimum: u32,
}

#[derive(serde::Deserialize, serde::Serialize, Clone, PartialEq)]
pub struct SrvRecord {
	priority: u16,
	weight: u16,
	port: u16,
	target: String,
}

/* Lookup Functions*/

pub async fn reverse_lookup(
		resolver: &TokioAsyncResolver,
		address: &IpAddr,
) -> Option<String> {
	let revese_res = resolver.reverse_lookup(*address);
	match revese_res.await {
		Ok(lookup) => {
			for name in lookup {
				return Some(name.to_string())
			}
			None
		}
		Err(e) => {
			let kind = e.kind();
			match kind {
				ResolveErrorKind::NoRecordsFound { .. } => {
					//Ignore, that just happens …
				}
				_ => {
					println!("Reverse lookup on {address} failed: {kind}");
				}
			}
			None
		}
	}	
}

pub fn opush<T: std::cmp::PartialEq>(opt_vec: &mut Option<Vec<T>>, data: T) {
	if opt_vec.is_none() {
		*opt_vec = Some(Default::default());
	}
	match opt_vec {
		Some(vec) => if !vec.contains(&data) { vec.push(data) },
		None => {},
	}
}

pub fn set_default_if_none<T>(opt_vec: &mut Option<Vec<T>>) {
	match opt_vec {
		Some(_) => {},
		None => {
			*opt_vec = Some(Default::default());
		}
	}
} 

pub fn add_record_to_lookup_result(result: &mut DnsLookupResult, record: &RData){
	match record {
		RData::AAAA(address) => opush(&mut result.aaaa, std::net::IpAddr::V6(*address)),
		RData::ANAME(aname) => opush(&mut result.aname, aname.to_string()),
		RData::A(address) => opush(&mut result.a, std::net::IpAddr::V4(*address)),
		RData::CAA(caa) => opush(&mut result.caa, caa.to_string()),
		RData::CNAME(cname) => opush(&mut result.cname, cname.to_string()),
		RData::MX(mx) => opush(&mut result.mx, MxRecord{
			preference: mx.preference(),
			exchange:   mx.exchange().to_string(),
		}),
		RData::NS(ns) => opush(&mut result.ns, ns.to_string()),
		RData::SOA(soa) => opush(&mut result.soa, SoaRecord{
			mname:   soa.mname().to_string(),
			rname:   soa.rname().to_string(),
			serial:  soa.serial(),
			refresh: soa.refresh(),
			retry:   soa.retry(),
			expire:  soa.expire(),
			minimum: soa.minimum(),
		}),
		RData::SRV(srv) => opush(&mut result.srv, SrvRecord{
			priority: srv.priority(),
			weight: srv.weight(),
			port: srv.port(),
			target: srv.target().to_string(),
		}),
		RData::TXT(txt) => {
			for text in txt.txt_data().iter() {
				opush(
					&mut result.txt,
					String::from_utf8_lossy(text).into_owned()
				);
			}
		},
		_ => { println!("Tried to add an unkown DNS record to results: {record}"); },
	}
}

pub fn integrate_lookup_result(dig_result: &mut DnsLookupResult, lookup_result: Result<Lookup, ResolveError>) {
	match lookup_result {
		Ok(lookup) => {
			match lookup.query().query_type() {
				RecordType::AAAA =>  set_default_if_none(&mut dig_result.aaaa),
				RecordType::ANAME => set_default_if_none(&mut dig_result.aname),
				RecordType::A =>     set_default_if_none(&mut dig_result.a),
				RecordType::CAA =>   set_default_if_none(&mut dig_result.caa),
				RecordType::CNAME => set_default_if_none(&mut dig_result.cname),
				RecordType::MX =>    set_default_if_none(&mut dig_result.mx),
				RecordType::NS =>    set_default_if_none(&mut dig_result.ns),
				RecordType::SOA =>   set_default_if_none(&mut dig_result.soa),
				RecordType::SRV =>   set_default_if_none(&mut dig_result.srv),
				RecordType::TXT =>   set_default_if_none(&mut dig_result.txt),
				_ => { /* This should not happen */ },
			};
			let name = lookup.query().name();
			for record in lookup.record_iter() {
				if name == record.name() {
					if let Some(data) = record.data() {
						add_record_to_lookup_result(dig_result, data);
					}
				}
				//TODO: handle additional responses
			}
		},
		Err(e) => {
			match e.kind() {
				ResolveErrorKind::Message(..) |
				ResolveErrorKind::Msg(..) |
				ResolveErrorKind::NoConnections |
				ResolveErrorKind::Io(..) |
				ResolveErrorKind::Proto(..) => {
					dig_result.other_error = true;
					println!("There was an error while doing a DNS Lookup: {e}");
				},
				ResolveErrorKind::Timeout => {
					dig_result.timeout = true;
					println!("There was a timeout while doing a DNS Lookup.");
				},
				ResolveErrorKind::NoRecordsFound{response_code, ..} => {
					match response_code {
						ResponseCode::NXDomain => dig_result.nxdomain = true,
						ResponseCode::NoError => {},
						_ => {
							println!("The DNS Server returned an error while doing a DNS Lookup: {response_code}");
							 dig_result.dns_error = true;
						},
					}
				}
				_ => { /*Ignore for now*/ },
			}
		}
	}
}

// This function takes a resolver, a domain name and returns a DnsLookupResult.
// If do_full_lookup is false only the A and AAAA (CNAMEs planned for the future)
// records will be fetched.
pub async fn lookup(
	resolver: &TokioAsyncResolver,
	name: &Name,
	do_full_lookup: bool,
) -> DnsLookupResult {
	let (
		ipv4_lookup_res,
		ipv6_lookup_res,
		cname_lookup_res,
		aname_lookup_res
	) = join!(
		resolver.lookup(name.clone(), RecordType::A),
		resolver.lookup(name.clone(), RecordType::AAAA),
		resolver.lookup(name.clone(), RecordType::CNAME),
		resolver.lookup(name.clone(), RecordType::ANAME),
	);

	// initlize an empty lookup result
	let mut dig_result: DnsLookupResult = Default::default();

	integrate_lookup_result(&mut dig_result, ipv4_lookup_res);
	integrate_lookup_result(&mut dig_result, ipv6_lookup_res);
	integrate_lookup_result(&mut dig_result, cname_lookup_res);
	integrate_lookup_result(&mut dig_result, aname_lookup_res);
	
	//Don't do an extented lookup if the domain seemingly doesn't exist
	if do_full_lookup && !dig_result.nxdomain {
		let (
			mx_lookup_res,
			ns_lookup_res,
			soa_lookup_res,
			caa_lookup_res,
			srv_lookup_res,
			txt_lookup_res
		) = join!(
			resolver.lookup(name.clone(), RecordType::MX),
			resolver.lookup(name.clone(), RecordType::NS),
			resolver.lookup(name.clone(), RecordType::SOA),
			resolver.lookup(name.clone(), RecordType::CAA),
			resolver.lookup(name.clone(), RecordType::SRV),
			resolver.lookup(name.clone(), RecordType::TXT),
		);

		integrate_lookup_result(&mut dig_result, mx_lookup_res);
		integrate_lookup_result(&mut dig_result, ns_lookup_res);
		integrate_lookup_result(&mut dig_result, soa_lookup_res);
		integrate_lookup_result(&mut dig_result, caa_lookup_res);
		integrate_lookup_result(&mut dig_result, srv_lookup_res);
		integrate_lookup_result(&mut dig_result, txt_lookup_res);

	}

	return dig_result

}