module Main (main)
where

import Control.Concurrent.ParallelIO.Global (
  parallelInterleaved,
  stopGlobalPool )
import Control.Monad (unless)
import qualified Data.ByteString.Char8 as BS (
  getContents,
  pack,
  words )
import Data.List ( (\\), intercalate, sort )
import Network.DNS (
  DNSError(NameError),
  Domain,
  FileOrNumericHost(RCHostName),
  ResolvConf(resolvInfo, resolvTimeout),
  defaultResolvConf,
  makeResolvSeed,
  normalize,
  normalizeCase,
  withResolver )
import System.Exit (ExitCode(ExitFailure), exitWith)
import System.IO (hPutStrLn, stderr)

import CommandLine (
   Args(NS,MX,delegates,no_append_root,server,timeout),
   get_args)
import DNS (
  LookupResult,
  lookupMX',
  lookupNS',
  resolve_address )
import ExitCodes (exit_bad_server)
import Timeout (Timeout(seconds))


-- | Given a list of delegates, report results for this
--   'LookupResult'.
--
--   If there's an empty list in the second component, there were no
--   query results, so we report that the domain was not delegated. If
--   there were some results and there are leftovers (after removing
--   the delegates), we report those as well.
--
--   Before processing, all names are normalized using the supplied
--   function @normalize_function@. Ideally this should be the same
--   function applied to the user-input names.
--
report :: [Domain] -- ^ The list of @delgts@
       -> (Domain -> Domain) -- ^ Domain name normalization function,
                             --   @normalize_function@.
       -> LookupResult
       -> IO ()

-- If the lookup succeeded but there were no results, report that the
-- domain is not delegated. Note that the behavior of the DNS library
-- changed with regard to this at some point: we used to get back
-- a "success," but with an empty list of results. Now a NameError
-- (which is not actually an error!) is returned.
report _ _ (d, Left NameError) =
  putStrLn $ "Domain " ++ (show d) ++ " not delegated."

-- If the lookup resulted in some other DNS error, we just ignore the
-- whole thing.
report _ _ (_, Left _) = return ()

-- Otherwise, subtract our delegates from the list of results and
-- report the leftovers.
report delgts normalize_function (d, Right raw_hosts) = do
  let nrml_hosts = map normalize_function raw_hosts
  -- Sort the leftovers so that we can test the expected output.
  let leftovers = sort (nrml_hosts \\ delgts)
  unless (null leftovers) $
    putStrLn $ "Domain " ++
               (show d) ++
               " delegates somewhere else: " ++
               (join " " (map show leftovers))
  where
    -- Create one big string by joining together a list of smaller
    -- strings and placing a delimiter between them.
    join :: String -> [String] -> String
    join delimiter strings = intercalate delimiter strings


main :: IO ()
main = do
  cfg <- get_args

  -- This reads stdin.
  input <- BS.getContents

  -- Split the input on any whitespace characters.
  let raw_domains = BS.words input

  -- Convert these to ByteStrings.
  let raw_delegates = map BS.pack (delegates cfg)

  let normalize_function =
        if (no_append_root cfg)
        then normalizeCase
        else normalize

  -- Normalize the given names and delegates
  let nrml_domains   = map normalize_function raw_domains
  let nrml_delegates = map normalize_function raw_delegates

  rc' <- case (server cfg) of
           Nothing -> return defaultResolvConf
           Just s -> do
             s' <- resolve_address s
             case s' of
               Left err -> do
                 let errmsg = show err
                 hPutStrLn stderr ("Bad DNS server or lookup error: " ++ errmsg)
                 exitWith (ExitFailure exit_bad_server)
               Right []  -> do
                 hPutStrLn stderr ("Hostname " ++ s ++ " has no 'A' records.")
                 exitWith (ExitFailure exit_bad_server)
               Right (srv:_) ->
                 return $ defaultResolvConf { resolvInfo =
                                                RCHostName (show srv) }

  -- Set the timeout from the command line. The resolvTimeout field is
  -- in microseconds, so we multiply by one million.
  let a_million = 1000 * 1000 :: Int
  let rc = rc' { resolvTimeout = a_million * seconds (timeout cfg) }
  rs <- makeResolvSeed rc

  let lookup_function = case cfg of
                          NS{} -> lookupNS'
                          MX{} -> lookupMX'

  let lookup' d = withResolver rs $ \resolver ->
                    lookup_function resolver d

  -- Construct a list of [IO whatever]. The withResolver calls
  -- are the ones that should be run in parallel.
  let actions = map lookup' nrml_domains

  -- Run the lookup actions in parallel.
  results <- parallelInterleaved actions

  -- Output the results.
  mapM_ (report nrml_delegates normalize_function) results

  stopGlobalPool
