module Network.DNS.Named
    ( runNamed
    )
    where

import           Control.Concurrent
import           Control.Exception
import           Control.Monad
import           Data.Binary
import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy as LBS
import           Data.Maybe
import           Network.Socket
import qualified Network.Socket.ByteString as NB
import           Network.DNS.Message
import           Network.DNS.Named.Config
import           Network.DNS.Named.Responder
import           Network.DNS.Named.ResponseBuilder
import           Network.DNS.Named.Zone
import           System.Posix.Signals


runNamed :: ZoneFinder zf => Config -> zf -> IO ()
runNamed cnf zf
    = withSocketsDo $
      do installHandler sigPIPE Ignore Nothing
         _tcpListenerTID <- forkIO $ tcpListen
         udpListen
    where
      udpListen :: IO ()
      udpListen = do -- FIXME: we should support IPv6 when the network package supports it.
                     so <- socket AF_INET Datagram defaultProtocol
                     print cnf
                     bindSocket so $ cnfServerAddress cnf
                     udpLoop so

      udpLoop :: Socket -> IO ()
      udpLoop so
          = do (packet, cameFrom) <- NB.recvFrom so 512
               _handlerTID <- forkIO $ udpHandler so packet cameFrom
               udpLoop so

      tcpListen :: IO ()
      tcpListen = putStrLn "FIXME: tcpListen is not implemented yet."

      udpHandler :: Socket -> BS.ByteString -> SockAddr -> IO ()
      udpHandler so packet cameFrom
          = do msg   <- evaluate $ unpackMessage packet
               msg'  <- handleMessage msg
                        `onException`
                        do let servfail = mkErrorReply ServerFailure msg
                           NB.sendTo so (packMessage (Just 512) servfail) cameFrom
               _sent <- NB.sendTo so (packMessage (Just 512) msg') cameFrom
               return ()

      handleMessage :: Message -> IO Message
      handleMessage msg
          = case validateQuery msg of
              NoError
                  -> do builders <- mapM handleQuestion $ msgQuestions msg

                        let builder = foldl (>>) (return ()) builders
                            msg'    = runBuilder msg builder

                        return msg'

              err -> return $ mkErrorReply err msg

      handleQuestion :: SomeQ -> IO (Builder ())
      handleQuestion (SomeQ q)
          = do zone    <- findZone zf (qName q)
               -- FIXME: this is merely a bogus implementation.
               -- It considers no additional or authoritative sections.
               results <- mapM (runResponder' q) (zoneResponders zone)
               return $ mapM_ addAnswer $ concat results


validateQuery :: Message -> ResponseCode
validateQuery = validateHeader . msgHeader
    where
      validateHeader :: Header -> ResponseCode
      validateHeader hdr
          | hdMessageType hdr /= Query         = NotImplemented
          | hdOpcode      hdr /= StandardQuery = NotImplemented
          | otherwise                          = NoError


packMessage :: Maybe Int -> Message -> BS.ByteString
packMessage limM = BS.concat . LBS.toChunks . truncateMsg
    where
      truncateMsg :: Message -> LBS.ByteString
      truncateMsg msg
          = let packet    = encode msg
                needTrunc = fromMaybe False $
                            do lim <- limM
                               return $ fromIntegral (LBS.length packet) > lim
            in
              if needTrunc then
                  truncateMsg $ trunc' msg
              else
                  packet

      trunc' :: Message -> Message
      trunc' msg
          | notNull $ msgAdditionals msg
              = msg {
                  msgAdditionals = truncList $ msgAdditionals msg
                }
          | notNull $ msgAuthorities msg
              = msg {
                  msgHeader      = setTruncFlag $ msgHeader msg
                , msgAuthorities = truncList $ msgAuthorities msg
                }
          | notNull $ msgAnswers msg
              = msg {
                  msgHeader      = setTruncFlag $ msgHeader msg
                , msgAnswers     = truncList $ msgAnswers msg
                }
          | notNull $ msgQuestions msg
              = msg {
                  msgHeader      = setTruncFlag $ msgHeader msg
                , msgQuestions   = truncList $ msgQuestions msg
                }
          | otherwise
              = error ("packMessage: You are already skinny and need no diet: " ++ show msg)

      setTruncFlag :: Header -> Header
      setTruncFlag hdr = hdr { hdIsTruncated = True }

      notNull :: [a] -> Bool
      notNull = not . null

      truncList :: [a] -> [a]
      truncList xs = take (length xs - 1) xs

unpackMessage :: BS.ByteString -> Message
unpackMessage = decode . LBS.fromChunks . return

mkErrorReply :: ResponseCode -> Message -> Message
mkErrorReply err msg
    = let header = msgHeader msg
          msg'   = msg {
                     msgHeader = header {
                                   hdMessageType           = Response
                                 , hdIsAuthoritativeAnswer = False
                                 , hdIsTruncated           = False
                                 , hdIsRecursionAvailable  = False
                                 , hdResponseCode          = err
                                 }
                   }
      in
        msg'