module DNS (
  MxSetMap,
  NormalDomain,
  mx_set_map,
  normalize_string )
where

import qualified Data.ByteString.Char8 as BS ( pack, unpack )
import Data.List ( nub )
import Data.Map ( Map )
import qualified Data.Map as Map ( fromList )
import Data.Set ( Set )
import qualified Data.Set as Set ( fromList )
import Network.DNS (
  Domain,
  defaultResolvConf,
  lookupMX,
  makeResolvSeed,
  normalize,
  withResolver )

-- | A type-safe wrapper around a domain name (represented as a
--   string) that ensures we've created it by calling
--   'normalize_string'. This prevents us from making
--   comparisons on un-normalized 'Domain's or 'String's.
--
newtype NormalDomain =
  NormalDomain String
  deriving ( Eq, Ord, Show )


-- | A set of mail exchanger names, represented as 'String's. The use
--   of 'NormalDomain' prevents us from constructing a set of names
--   that aren't normalized first.
--
type MxSet = Set NormalDomain


-- | A map from domain names (represented as 'String's) to sets of
--   mail exchanger names (also represented as 'String's). The use of
--   'NormalDomain' in the key prevents us from using keys that aren't
--   normalized; this is important because we'll be using them for
--   lookups and want e.g. \"foo.com\" and \"FOO.com\" to look up the
--   same MX records.
--
type MxSetMap = Map NormalDomain MxSet


-- | Normalize a domain name string by converting to a 'Domain',
--   calling 'normalize', and then converting back.
--
--   ==== __Examples__
--
--   >>> normalize_string "ExAMplE.com"
--   NormalDomain "example.com."
--
normalize_string :: String -> NormalDomain
normalize_string = NormalDomain . BS.unpack . normalize . BS.pack


-- | Retrieve all MX records for the given domain. This is somewhat
--   inefficient, since we create the resolver every time.
--
lookup_mxs :: Domain -> IO [Domain]
lookup_mxs domain = do
  default_rs <- makeResolvSeed defaultResolvConf
  withResolver default_rs $ \resolver -> do
    mxs <- lookupMX resolver domain
    return $ case mxs of
               Left  _     -> []
               Right pairs -> map fst pairs


-- | Takes a list of domain names represented as 'String's and
--   constructs a map from domain names to sets of mail exchangers
--   (for those domain names) also represented as 'String's.
--
--   During construction, we have to switch to the DNS internal
--   representation of a 'Domain' which uses ByteStrings, but before
--   we return the map to the client, we want everything to be in
--   terms of standard 'String's for comparison purposes.
--
--   The list of domains is normalized and de-duped before lookups are
--   performed to avoid doing lookups twice for identical domains.
--
mx_set_map :: [String] -> IO MxSetMap
mx_set_map domains = do
  -- Construct a list of pairs.
  pairs <- mapM make_pair unique_domains

  -- And make a map from the pairs.
  return $ Map.fromList pairs

  where
    -- Convert, normalize, and de-dupe the @domains@.
    unique_domains :: [Domain]
    unique_domains = nub $ map (normalize . BS.pack) domains

    -- | Convert a string domain name into a pair containing the
    --   domain name in the first component and a set of its mail
    --   exchangers (as strings) in the second component.
    --
    make_pair :: Domain -> IO (NormalDomain, Set NormalDomain)
    make_pair domain = do
      -- Lookup the @domain@'s MX records.
      mx_list <- lookup_mxs domain

      -- Now convert the MX records *back* to strings, and then to
      -- NormalDomains
      let normal_mx_list = map (normalize_string . BS.unpack) mx_list

      -- Convert the list into a set...
      let normal_mx_set = Set.fromList normal_mx_list

      -- The lookup key.
      let normal_domain = normalize_string $ BS.unpack domain

      -- Finally, construct the pair and return it.
      return (normal_domain, normal_mx_set)