    {-# LANGUAGE RecordWildCards #-}

    module Network.DNS.Encode (
        encode
      , composeQuery
      ) where

    import qualified Blaze.ByteString.Builder as BB (toByteString, fromWrite, writeInt16be)
    import qualified Data.ByteString.Lazy.Char8 as BL (ByteString)
    import qualified Data.ByteString.Char8 as BS (length, null, break, drop)
    import Network.DNS.StateBinary
    import Network.DNS.Internal
    import Data.Monoid
    import Control.Monad.State
    import Data.Bits
    import Data.Word
    import Data.IP

    (+++) :: Monoid a => a -> a -> a
    (+++) = mappend

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

    -- | Composing query. First argument is a number to identify response.

    composeQuery :: Int -> [Question] -> BL.ByteString
    composeQuery idt qs = encode qry
      where
        hdr = header defaultQuery
        qry = defaultQuery {
            header = hdr {
               identifier = idt
             , qdCount = length qs
             }
          , question = qs
          }

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

    -- | Composing DNS data.

    encode :: DNSFormat -> BL.ByteString
    encode fmt = runSPut (encodeDNSFormat fmt)

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

    encodeDNSFormat :: DNSFormat -> SPut
    encodeDNSFormat fmt = encodeHeader hdr
                      +++ mconcat (map encodeQuestion qs)
                      +++ mconcat (map encodeRR an)
                      +++ mconcat (map encodeRR au)
                      +++ mconcat (map encodeRR ad)
      where
        hdr = header fmt
        qs = question fmt
        an = answer fmt
        au = authority fmt
        ad = additional fmt

    encodeHeader :: DNSHeader -> SPut
    encodeHeader hdr = encodeIdentifier (identifier hdr)
                   +++ encodeFlags (flags hdr)
                   +++ decodeQdCount (qdCount hdr)
                   +++ decodeAnCount (anCount hdr)
                   +++ decodeNsCount (nsCount hdr)
                   +++ decodeArCount (arCount hdr)
      where
        encodeIdentifier = putInt16
        decodeQdCount = putInt16
        decodeAnCount = putInt16
        decodeNsCount = putInt16
        decodeArCount = putInt16

    encodeFlags :: DNSFlags -> SPut
    encodeFlags DNSFlags{..} = put16 word
      where
        word16 :: Enum a => a -> Word16
        word16 = toEnum . fromEnum

        set :: Word16 -> State Word16 ()
        set byte = modify (.|. byte)

        st :: State Word16 ()
        st = sequence_
                  [ set (word16 rcode)
                  , when recAvailable        $ set (bit 7)
                  , when recDesired          $ set (bit 8)
                  , when trunCation          $ set (bit 9)
                  , when authAnswer          $ set (bit 10)
                  , set (word16 opcode `shiftL` 11)
                  , when (qOrR==QR_Response) $ set (bit 15)
                  ]

        word = execState st 0

    encodeQuestion :: Question -> SPut
    encodeQuestion Question{..} =
            encodeDomain qname
        +++ putInt16 (typeToInt qtype)
        +++ put16 1

    encodeRR :: ResourceRecord -> SPut
    encodeRR ResourceRecord{..} =
        mconcat
          [ encodeDomain rrname
          , putInt16 (typeToInt rrtype)
          , put16 1
          , putInt32 rrttl
          , rlenRDATA
          ]
      where
        -- Encoding rdata without using rdlen
        rlenRDATA = do
            addPositionW 2 -- "simulate" putInt16
            rDataWrite <- encodeRDATA rdata
            let rdataLength = fromIntegral . BS.length . BB.toByteString . BB.fromWrite $ rDataWrite
            let rlenWrite = BB.writeInt16be rdataLength
            return rlenWrite +++ return rDataWrite

    encodeRDATA :: RDATA -> SPut
    encodeRDATA rd = case rd of
        (RD_A ip)          -> mconcat $ map putInt8 (fromIPv4 ip)
        (RD_AAAA ip)       -> mconcat $ map putInt16 (fromIPv6 ip)
        (RD_NS dom)        -> encodeDomain dom
        (RD_CNAME dom)     -> encodeDomain dom
        (RD_PTR dom)       -> encodeDomain dom
        (RD_MX prf dom)    -> mconcat [putInt16 prf, encodeDomain dom]
        (RD_TXT txt)       -> putByteString txt
        (RD_OTH bytes)     -> mconcat $ map putInt8 bytes
        (RD_SOA d1 d2 serial refresh retry expire min') -> mconcat
            [ encodeDomain d1
            , encodeDomain d2
            , putInt32 serial
            , putInt32 refresh
            , putInt32 retry
            , putInt32 expire
            , putInt32 min'
            ]
        (RD_SRV prio weight port dom) -> mconcat
            [ putInt16 prio
            , putInt16 weight
            , putInt16 port
            , encodeDomain dom
            ]

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

    encodeDomain :: Domain -> SPut
    encodeDomain dom | BS.null dom = put8 0
    encodeDomain dom = do
        mpos <- wsPop dom
        cur <- gets wsPosition
        case mpos of
            Just pos -> encodePointer pos
            Nothing  -> wsPush dom cur >>
                        mconcat [ encodePartialDomain hd
                                , encodeDomain tl
                                ]
      where
        (hd, tl') = BS.break (=='.') dom
        tl = if BS.null tl' then tl' else BS.drop 1 tl'

    encodePointer :: Int -> SPut
    encodePointer pos = let w = (pos .|. 0xc000) in putInt16 w

    encodePartialDomain :: Domain -> SPut
    encodePartialDomain sub = putInt8 (BS.length sub)
                          +++ putByteString sub