use crate::trie_map::TrieMap;
use crate::util::OwnedRecordData;
use domain::base::{Dname, Question, Record, Rtype};
use domain::rdata::{Aaaa, AllRecordData, A};
use lazy_static::lazy_static;
use std::collections::{HashMap, HashSet};
use std::net::{IpAddr, Ipv4Addr};

lazy_static! {
    // Put a simple blocklist of domains at ../blocklist.txt
    // All domains in the file will be resolved to 0.0.0.0
    // This can be used for ad-blocking, as converting the
    // blocklists to JSON config file would not be a great idea,
    // but converting them to a dumb list of domains should be trivial
    static ref BLOCK_LIST: HashSet<String> = parse_blocklist_file();
}

fn parse_blocklist_file() -> HashSet<String> {
    let mut ret = HashSet::new();
    for line in include_str!("../blocklist.txt").lines() {
        if line.is_empty() {
            continue;
        }

        if line.starts_with("#") {
            continue;
        }

        ret.insert(line.trim().to_string());
    }
    ret
}

pub struct OverrideResolver {
    simple_matches: HashMap<String, IpAddr>,
    suffix_matches: TrieMap<IpAddr>,
    override_ttl: u32,
}

impl OverrideResolver {
    pub fn new(overrides: HashMap<String, String>, override_ttl: u32) -> OverrideResolver {
        let (simple_matches, suffix_matches) = Self::build_match_tables(overrides);
        OverrideResolver {
            suffix_matches,
            simple_matches,
            override_ttl,
        }
    }

    fn build_match_tables(
        overrides: HashMap<String, String>,
    ) -> (HashMap<String, IpAddr>, TrieMap<IpAddr>) {
        let mut simple = HashMap::new();
        let mut suffix = TrieMap::new();
        for (k, v) in overrides.into_iter() {
            match v.parse::<IpAddr>() {
                Ok(addr) => {
                    if k.starts_with("*.") {
                        // Anything starting with a wildcard character is a suffix match
                        // we convert it to a prefix match by reversing the domain
                        // Note that we get rid of the wildcard but keep the dot, i.e.
                        // we don't allow suffix match in the middle of a part of a domain
                        suffix.put_prefix(k[1..].chars().rev().collect::<String>(), addr);
                    } else {
                        simple.insert(k, addr);
                    }
                }
                // Ignore malformed IP addresses
                Err(_) => continue,
            }
        }
        (simple, suffix)
    }

    pub fn try_resolve(
        &self,
        question: &Question<Dname<Vec<u8>>>,
    ) -> Option<Record<Dname<Vec<u8>>, OwnedRecordData>> {
        match question.qtype() {
            // We only handle resolution of IP addresses
            Rtype::A | Rtype::A6 | Rtype::Aaaa | Rtype::Cname | Rtype::Any => (),
            // So if the question is anything else, just skip
            _ => return None,
        }

        let name = question.qname().to_string();
        if let Some(addr) = self.simple_matches.get(&name) {
            self.respond_with_addr(question, addr)
        } else if BLOCK_LIST.get(&name).is_some() {
            self.respond_with_addr(question, &IpAddr::V4(Ipv4Addr::UNSPECIFIED))
        } else if let Some(addr) = self
            .suffix_matches
            .get_by_prefix(name.chars().rev().collect::<String>())
        {
            self.respond_with_addr(question, addr)
        } else {
            None
        }
    }

    fn respond_with_addr(
        &self,
        question: &Question<Dname<Vec<u8>>>,
        addr: &IpAddr,
    ) -> Option<Record<Dname<Vec<u8>>, OwnedRecordData>> {
        let rdata: OwnedRecordData = match addr {
            IpAddr::V4(addr) => AllRecordData::A(A::new(addr.clone())),
            IpAddr::V6(addr) => AllRecordData::Aaaa(Aaaa::new(addr.clone())),
        };

        let record = Record::new(
            question.qname().clone(),
            question.qclass(),
            self.override_ttl,
            rdata,
        );
        return Some(record);
    }
}