{-# LANGUAGE CApiFFI #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DeriveFoldable             #-}
{-# LANGUAGE DeriveFunctor              #-}
{-# LANGUAGE DeriveTraversable          #-}

-- |
-- Copyright: © 2017 Herbert Valerio Riedel
-- License: GPLv3
--
module Network.DNS.FFI where

import           Control.Applicative as App
import           Control.DeepSeq
import           Control.Exception
import           Control.Monad
import           Data.Bits
import           Data.ByteString (ByteString)
import qualified Data.ByteString as BS
import           Data.Int
import           Data.Word
import           Foreign.C.Types
import           Foreign.Marshal.Alloc
import           Foreign.Ptr
import           Foreign.Storable
import           Numeric (showHex)

#include <windows.h>
#include <windns.h>

type DWORD = #type DWORD
type WORD  = #type WORD
type BOOL  = #type BOOL

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

-- | Cache time-to-live expressed in seconds
newtype TTL  = TTL Int32 deriving (Eq,Ord,Read,Show,NFData)

-- | @\<character-string\>@ as per [RFC 1035, section 3.3](https://tools.ietf.org/html/rfc1035#section-3.3).
--
-- A sequence of up to 255 octets
--
-- The limit of 255 octets is caused by the encoding which uses by a
-- prefixed octet denoting the length.
newtype CharStr = CharStr ByteString deriving (Eq,Ord,Read,Show)

instance NFData CharStr where rnf (CharStr !_) = ()

-- | @\<domain-name\>@ as per [RFC 1035, section 3.3](https://tools.ietf.org/html/rfc1035#section-3.3).
--
-- A domain-name represented as a series of labels separated by dots.
newtype Name    = Name ByteString deriving (Eq,Ord,Read,Show)

instance NFData Name where rnf (Name !_) = ()


-- | An IPv6 address
--
-- The IP address is represented in network order,
-- i.e. @2606:2800:220:1:248:1893:25c8:1946@ is
-- represented as @(IPv6 0x2606280002200001 0x248189325c81946)@.
data IPv6 = IPv6 !Word64 !Word64
          deriving (Eq,Ord,Read)

instance Show IPv6 where
    showsPrec p (IPv6 hi lo) = showParen (p >= 11) (showString "IPv6 0x" . showHex hi . showString " 0x" . showHex lo)

instance NFData IPv6 where
  rnf (IPv6 _ _) = ()

mkIPv6 :: Word32 -> Word32 -> Word32 -> Word32 -> IPv6
mkIPv6 a b c d = IPv6 (mkW64 a b) (mkW64 c d)
  where
    mkW64 x y = (fromIntegral (byteSwap32 x) `shiftL` 32) .|. fromIntegral (byteSwap32 y) :: Word64

-- | An IPv4 address
--
-- The IP address is represented in network order, i.e. @127.0.0.1@ is
-- represented as @(IPv4 0x7f000001)@.
data IPv4 = IPv4 !Word32
          deriving (Eq,Ord,Read)

instance NFData IPv4 where
  rnf (IPv4 _) = ()

mkIPv4 :: Word32 -> IPv4
mkIPv4 a = IPv4 (byteSwap32 a)

instance Show IPv4 where
    showsPrec p (IPv4 n) = showParen (p >= 11) (showString "IPv4 0x" . showHex n)

-- | @SRV@ Record data as per [RFC 2782](https://tools.ietf.org/html/rfc2782)
data SRV l = SRV { srvPriority :: !Word16
                 , srvWeight   :: !Word16
                 , srvPort     :: !Word16
                 , srvTarget   :: !l
                 } deriving (Eq,Read,Show,Functor,Foldable,Traversable)

instance NFData l => NFData (SRV l) where
  rnf (SRV _ _ _ l) = rnf l

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

dnsQuery :: Bool -> Name -> DnsType -> IO (Either Int [DnsRecord])
dnsQuery exact (Name n) ty = do
  alloca $ \pst -> do
    BS.useAsCString n $ \n' ->
      bracket (c_dns_query n' (fromDnsType ty) pst)
              (c_free_record)
              $ \p0 -> do
                st <- peek pst
                if (st /= 0)
                  then pure (Left (fromIntegral st))
                  else do
                    tmp <- travRecs (peekRec n') p0
                    if exact
                      then pure (Right [ r | (b,r) <- tmp, b ])
                      else pure (Right (map snd tmp))


foreign import capi safe "hs_windns.h hs_dns_query" c_dns_query :: Ptr CChar -> WORD -> Ptr CLong -> IO (Ptr DnsRecord)

foreign import capi unsafe "hs_windns.h hs_free_record" c_free_record :: Ptr DnsRecord -> IO (Ptr DnsRecord)

foreign import capi unsafe "hs_windns.h DnsNameCompare_A" c_dns_name_eq :: Ptr CChar -> Ptr CChar -> IO BOOL


travRecs :: (Ptr DnsRecord -> IO a) -> Ptr DnsRecord -> IO [a]
travRecs f p0 = go [] p0
  where
    go acc p
      | p == nullPtr = App.pure (reverse acc)
      | otherwise    = do
          x <- f p
          p' <- next p
          go (x:acc) p'

    next :: Ptr DnsRecord -> IO (Ptr DnsRecord)
    next = #{peek DNS_RECORDA, pNext}


peekRec :: Ptr CChar -> Ptr DnsRecord -> IO (Bool,DnsRecord)
peekRec n0 p = do
  drNamePtr <- #{peek DNS_RECORDA, pName} p
  same <- c_dns_name_eq n0 drNamePtr
  drName <- Name <$> BS.packCString drNamePtr
  drType <- toDnsType <$> #{peek DNS_RECORDA, wType} p
  drTTL  <- TTL . fromIntegral <$> (#{peek DNS_RECORDA, dwTtl} p :: IO DWORD)

  drData <- case drType of
              DnsTypeA    -> DnsDataA . mkIPv4 <$> #{peek DNS_RECORDA, Data.A.IpAddress} p
              DnsTypeAAAA -> DnsDataAAAA <$> (mkIPv6 <$> #{peek DNS_RECORDA, Data.AAAA.Ip6Address.IP6Dword[0]} p
                                                     <*> #{peek DNS_RECORDA, Data.AAAA.Ip6Address.IP6Dword[1]} p
                                                     <*> #{peek DNS_RECORDA, Data.AAAA.Ip6Address.IP6Dword[2]} p
                                                     <*> #{peek DNS_RECORDA, Data.AAAA.Ip6Address.IP6Dword[3]} p)
              DnsTypeTXT  -> do
                cnt <- #{peek DNS_RECORDA, Data.TXT.dwStringCount} p
                let ptr0 = #{ptr DNS_RECORDA, Data.TXT.pStringArray[0]} p
                tptrs <- forM [0.. fromIntegral (cnt :: DWORD)-1] (peekElemOff ptr0)
                DnsDataTXT <$> mapM (fmap CharStr . BS.packCString) tptrs

              DnsTypeCNAME -> do
                cnamptr <- #{peek DNS_RECORDA, Data.CNAME.pNameHost} p
                DnsDataCNAME . Name <$> BS.packCString cnamptr

              DnsTypeSRV -> do
                namptr      <- #{peek DNS_RECORDA, Data.SRV.pNameTarget} p
                srvTarget   <- Name <$> BS.packCString namptr
                srvPort     <- #{peek DNS_RECORDA, Data.SRV.wPort} p
                srvWeight   <- #{peek DNS_RECORDA, Data.SRV.wWeight} p
                srvPriority <- #{peek DNS_RECORDA, Data.SRV.wPriority} p

                pure $! DnsDataSRV (SRV {..})

              DnsType w -> pure (DnsData w)

  evaluate $ force (same /= 0,DnsRecord{..})

data DnsRecord = DnsRecord
    { drName :: !Name
    , drType :: !DnsType
    , drTTL  :: !TTL
    , drData :: !DnsData
    } deriving Show

instance NFData DnsRecord where
  rnf (DnsRecord n y t d) = n `deepseq` y `deepseq` t `deepseq` d `deepseq` ()


data DnsData = DnsDataA    !IPv4
             | DnsDataAAAA !IPv6
             | DnsDataTXT  [CharStr]
             | DnsDataCNAME !Name
             | DnsDataSRV  !(SRV Name)
             | DnsData     !WORD -- unknown/unsupported
             deriving Show

instance NFData DnsData where
  rnf (DnsDataA _) = ()
  rnf (DnsDataAAAA {}) = ()
  rnf (DnsDataTXT ts) = rnf ts
  rnf (DnsDataCNAME n) = rnf n
  rnf (DnsDataSRV srv) = rnf srv
  rnf (DnsData _) = ()

dnsDataType :: DnsData -> DnsType
dnsDataType DnsDataA {}     = DnsTypeA
dnsDataType DnsDataAAAA {}  = DnsTypeAAAA
dnsDataType DnsDataTXT {}   = DnsTypeTXT
dnsDataType DnsDataCNAME {} = DnsTypeCNAME
dnsDataType DnsDataSRV {}   = DnsTypeSRV
dnsDataType (DnsData w)     = DnsType w

data DnsType = DnsTypeA
             | DnsTypeAAAA
             | DnsTypeTXT
             | DnsTypeCNAME
             | DnsTypeSRV
             | DnsType !WORD
             deriving (Show)

instance NFData DnsType where rnf t = seq t ()

eqType :: DnsType -> DnsType -> Bool
eqType x y = fromDnsType x == fromDnsType y

fromDnsType :: DnsType -> WORD
fromDnsType x = case x of
  DnsTypeA     -> #const DNS_TYPE_A
  DnsTypeAAAA  -> #const DNS_TYPE_AAAA
  DnsTypeTXT   -> #const DNS_TYPE_TEXT
  DnsTypeCNAME -> #const DNS_TYPE_CNAME
  DnsTypeSRV   -> #const DNS_TYPE_SRV
  DnsType w    -> w

toDnsType :: WORD -> DnsType
toDnsType w = case w of
  #{const DNS_TYPE_A   }  -> DnsTypeA
  #{const DNS_TYPE_AAAA}  -> DnsTypeAAAA
  #{const DNS_TYPE_TEXT}  -> DnsTypeTXT
  #{const DNS_TYPE_CNAME} -> DnsTypeCNAME
  #{const DNS_TYPE_SRV}   -> DnsTypeSRV
  _                       -> DnsType w


#if !MIN_VERSION_base(4,7,0)
byteSwap32 :: Word32 -> Word32
byteSwap32 x = ((x .&. 0x000000FF) `shiftL` 24) .|.
               ((x .&. 0x0000FF00) `shiftL`  8) .|.
               ((x .&. 0x00FF0000) `shiftR`  8) .|.
               ((x .&. 0xFF000000) `shiftR` 24);
#endif