{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Main where

import Control.Concurrent
import Control.Monad
import Data.Either
import Data.List
import Data.Foldable
import Network.DNS.Lookup
import Network.DNS.Resolver
import Network.DNS.Types
import Network.Whois
import System.Posix.Files
import qualified Data.ByteString.Char8 as T
import qualified Data.Map as M
import Control.Monad.Trans.Maybe
import System.Environment

-- config
config_dns1     = "1.1.1.1"     -- cloudfare DNS
config_dns2     = "8.8.8.8"     -- google DNS
config_sleep    = 300000        -- thread sleep 300ms

config_tld      = "tld.txt"     -- top level domains to check
config_cache    = "cache.txt"   -- cache answers

-- check if given domain name has a whois entry
isDomainWhois :: String -> IO Bool
isDomainWhois x = check <$> whois x
                  where check (_,Nothing) = False
                        check _ = True

-- check if given domain has a dns entry (pass ResolvSeed as first argument)
isDomainDNS :: ResolvSeed -> String -> IO Bool
isDomainDNS rs domain = do a <- withResolver rs $ \resolver -> 
                                lookupA resolver $ T.pack domain
                           return $ either (not.isNameError) (const True) a
                           where isNameError NameError = True
                                 isNameError OperationRefused = error "operation refused"
                                 isNameError TimeoutExpired = error "timeout"
--                               isNameError ServerFailure = error "server failure"
                                 isNameError _ = False

-- group our domain name checkers in an Array (dns1, dns2, whois)
isDomain :: IO [String-> IO Bool]
isDomain = sequence [ 
                      isDomainDNS <$> makeResolvSeed defaultResolvConf 
                                    {resolvInfo = RCHostName config_dns1}
                     ,isDomainDNS <$> makeResolvSeed defaultResolvConf 
                                    {resolvInfo = RCHostName config_dns2}
--                     ,return isDomainWhois
                    ]

-- assure the domain name is free using all the funcs from the array
multiCheck :: [String -> IO Bool] -> String -> IO Bool
multiCheck funcs dom = do res <- runner $ map ($dom) funcs
                          return $ maybe False (const True) res
                          where runner=runMaybeT . asum . fmap (MaybeT . fmap guard)

-- load a Monoid from File (or init File with an empty Monoid)
loadFromFile :: forall a. (Show a, Monoid a,Read a) => FilePath -> IO a
loadFromFile cf = do fileExist cf >>= \x -> when (not x) $ 
                                writeFile cf (show (mempty::a))
                     read . T.unpack <$> T.readFile cf

-- wait until the length of the mvar array reaches given length,
-- so we now processing has finished.
waitFinished :: MVar [a] -> Int -> IO ()
waitFinished mvout l = readMVar mvout >>= \o -> when (length o /= l) $ 
                       do threadDelay config_sleep
                          -- putStrLn $ show (length o) ++ "/" ++ show l
                          waitFinished mvout l


-- given an array of testing functions a list of domains, output array 
-- and cache map. check if domains are free. for now we just use first 
-- function for checking! consider different approach
domainChecker :: [String -> IO Bool] -> MVar [String] -> MVar [(Bool,String)] 
                         -> MVar (M.Map String Bool) -> IO ()
domainChecker chk mvin mvout cache = 
            nxtDom >>= \dom ->  when (dom/="") $ do
                    c<-readMVar cache             
                    (r,ch)<-case M.lookup dom c of 
                        Just xx-> return (xx,True)
                        Nothing-> do x<-multiCheck chk dom 
                                     addCache x dom
                                     return (x,False)

                    outDom r dom
                    putStrLn $ dom  ++ " " 
                                    ++ (if r then "[\ESC[31m\STXTaken\ESC[m\STX]" else "[\ESC[32m\STXFree\ESC[m\STX]") 
                                    ++ (if ch then " (\ESC[34m\STXcache hit\ESC[m\STX)" else "") 
                    domainChecker chk mvin mvout cache
                    
                where outDom res dom =modifyMVar mvout $ \d -> 
                                          return ((res,dom):d,())

                      nxtDom         =modifyMVar mvin $ \d -> 
                                          if null d then return (d,"") 
                                          else return (tail d,head d)
                      addCache r dom =modifyMVar cache $ \d ->
                                          return (M.insert dom r d,())  

---------------------------------------------

-- main / entry
main :: IO()
--main = test
main = entry

-- testing/debugging
test :: IO ()
test = do 
    let dom="shememem.com"
    rs <- makeResolvSeed defaultResolvConf {resolvInfo = RCHostName config_dns1}
    a  <- withResolver rs $ \resolver -> lookupA resolver $ T.pack dom    
    print a
    x <- isDomainDNS rs dom
    print x

entry :: IO ()
entry = do 

          -- get cli params
          args<-getArgs
          progname<-getProgName
          when (length args<2) $ error $ "Usage: "++progname++" [domains_filename] [number_of_threads]"
          let config_domains= args !!0 
          let config_threads= read $ args !!1 

          -- read cache, domains, and tld from files 
          -- cache <- loadFromFile config_cache
          cache <- loadFromFile config_cache
          dom   <- readFile config_domains
          tld   <- readFile config_tld

          -- combine domains and tlds
          let domains=liftM2 (++) (uncomment $ lines dom) $ 
                                  map ("."++) (uncomment $ lines tld)

          -- init mvars
          mvarCache <- newMVar cache
          mvarIn    <- newMVar domains
          mvarOut   <- newMVar []

          -- check domains (in multiple threads)
          checkers <- isDomain
          sequence_ $ replicate config_threads $ forkIO $ 
                             domainChecker checkers mvarIn mvarOut mvarCache

          -- wait until all threads finish and get result
          waitFinished mvarOut (length domains)
          res <- readMVar mvarOut

          -- show free domains
          putStrLn "--- Free Domains ----"
          putStr $ unlines $ sort $ map snd $ filter (not.fst) res

          -- persist updated cache
          cache'<-takeMVar mvarCache 
          writeFile config_cache $ show cache'
          
          where uncomment=map (filter (/='\r')) . filter ((/='-').head)