'''
Shared load/save methods for intermediate results
'''

import functools
import logging
import pickle
from typing import Dict, Set

import dns.name
from ednsevalzone import AnIPAddress

def load_nsname2ipset() -> Dict[dns.name.Name, Set[AnIPAddress]]:
    """raises FileNotFoundError"""
    logging.info('loading NS name -> IP address mapping')
    with open('nsname2ipset.pickle', 'rb') as nsname2ipset_pickle:
        nsname2ipset = pickle.load(nsname2ipset_pickle)
    ip_cnt = functools.reduce(lambda cnt, ipset: cnt + len(ipset), nsname2ipset.values(), 0)
    logging.info('loaded %s unique IP addresses for %s NS names',
                 ip_cnt, len(nsname2ipset))
    return nsname2ipset

def load_domain2ipset() -> Dict[dns.name.Name, Set[AnIPAddress]]:
    """raises FileNotFoundError"""
    logging.info('loading domain-IP sets')
    with open('domain2ipset.pickle', 'rb') as domain2ipset_pickle:
        domain2ipset = pickle.load(domain2ipset_pickle)
    return domain2ipset

def save_nsname2ipset(nsname2ipset: Dict[dns.name.Name, Set[AnIPAddress]]) -> None:
    if nsname2ipset:
        ip_cnt = functools.reduce(lambda cnt, ipset: cnt + len(ipset), nsname2ipset.values(), 0)
    else:
        ip_cnt = 0
    logging.info('pickling NS name -> IP address mapping '
                 '(%s unique IP addresses for %s NS names)',
                 ip_cnt, len(nsname2ipset))
    with open('nsname2ipset.pickle', 'wb') as nsname2ipset_pickle:
        pickle.dump(
            nsname2ipset,
            nsname2ipset_pickle
            )

class DataStore():
    def __init__(self):
        self.zone = None
        self.nsname2ipset = None