module Network.DNS.Message
    ( Message(..)
    , MessageType(..)
    , Header(..)
    , Opcode(..)
    , ResponseCode(..)
    , Question(..)
    , ResourceRecord(..)
    , RecordType(..)
    , RecordClass(..)
    )
    where

import           Data.Binary
import           Data.Binary.Get
import           Data.Binary.Put
import           Data.Bits
import           Data.Word


data Message
    = Message {
        msgHeader      :: !Header
      , msgQuestions   :: ![Question]
      , msgAnswers     :: ![ResourceRecord]
      , msgAuthorities :: ![ResourceRecord]
      , msgAdditionals :: ![ResourceRecord]
      }

data Header
    = Header {
        hdMessageID             :: !Word16
      , hdMessageType           :: !MessageType
      , hdOpcode                :: !Opcode
      , hdIsAuthoritativeAnswer :: !Bool
      , hdIsTruncated           :: !Bool
      , hdIsRecursionDesired    :: !Bool
      , hdIsRecursionAvailable  :: !Bool
      , hdResponseCode          :: !ResponseCode

      -- These fields are supressed in this data structure:
      -- * QDCOUNT
      -- * ANCOUNT
      -- * NSCOUNT
      -- * ARCOUNT
      }

data MessageType
    = Query
    | Response
    deriving (Show, Eq)

data Opcode
    = StandardQuery
    | InverseQuery
    | ServerStatusRequest
    deriving (Show, Eq)

data ResponseCode
    = NoError
    | FormatError
    | ServerFailure
    | NameError
    | NotImplemented
    | Refused
    deriving (Show, Eq)

data Question
    = Question {
        qName  :: !DomainName
      , qType  :: !RecordType
      , qClass :: !RecordClass
      }
    deriving (Show, Eq)

type DomainName = [[Word8]]

data RecordClass
    = IN
    | CS -- Obsolete
    | CH
    | HS
    | AnyClass -- Only for queries
    deriving (Show, Eq)

data ResourceRecord
    = ResourceRecord {
        rrName  :: !DomainName
      , rrType  :: !RecordType
      , rrClass :: !RecordClass
      , rrTTL   :: !Word32
      , rrData  :: ![Word8]
      }
    deriving (Show, Eq)

data RecordType
    = A
    | NS
    | MD
    | MF
    | CNAME
    | SOA
    | MB
    | MG
    | MR
    | NULL
    | WKS
    | PTR
    | HINFO
    | MINFO
    | MX
    | TXT

    -- Only for queries:
    | AXFR
    | MAILB -- Obsolete
    | MAILA -- Obsolete
    | AnyType
    deriving (Show, Eq)

instance Binary Header where
    put h = do putWord16be $ hdMessageID h
               let qr    = fromIntegral $ fromEnum $ hdMessageType h
                   op    = fromIntegral $ fromEnum $ hdOpcode h
                   aa    = boolToNum $ hdIsAuthoritativeAnswer h
                   tc    = boolToNum $ hdIsTruncated h
                   rd    = boolToNum $ hdIsRecursionDesired h
                   ra    = boolToNum $ hdIsRecursionAvailable h
                   rc    = fromIntegral $ fromEnum $ hdResponseCode h
                   flags = ((qr `shiftL` 15) .&. 0x01) .|.
                           ((op `shiftL` 11) .&. 0x0F) .|.
                           ((aa `shiftL` 10) .&. 0x01) .|.
                           ((tc `shiftL`  9) .&. 0x01) .|.
                           ((rd `shiftL`  8) .&. 0x01) .|.
                           ((ra `shiftL`  7) .&. 0x01) .|.
                           ((rc `shiftL`  0) .&. 0x0F)
               putWord16be flags
        where
          boolToNum :: Num a => Bool -> a
          boolToNum True  = 1
          boolToNum False = 0

    get = do mID   <- getWord16be
             flags <- getWord16be
             let qr = toEnum $ fromIntegral ((flags `shiftR` 15) .&. 0x01)
                 op = toEnum $ fromIntegral ((flags `shiftR` 11) .&. 0x0F)
                 aa = testBit flags 10
                 tc = testBit flags 9
                 rd = testBit flags 8
                 ra = testBit flags 7
                 rc = toEnum $ fromIntegral (flags .&. 0x0F)
                 hd = Header {
                         hdMessageID             = mID
                       , hdMessageType           = qr
                       , hdOpcode                = op
                       , hdIsAuthoritativeAnswer = aa
                       , hdIsTruncated           = tc
                       , hdIsRecursionDesired    = rd
                       , hdIsRecursionAvailable  = ra
                       , hdResponseCode          = rc
                       }
             return hd

instance Enum MessageType where
    fromEnum Query    = 0
    fromEnum Response = 1

    toEnum 0 = Query
    toEnum 1 = Response
    toEnum _ = undefined

instance Enum Opcode where
    fromEnum StandardQuery       = 0
    fromEnum InverseQuery        = 1
    fromEnum ServerStatusRequest = 2

    toEnum 0 = StandardQuery
    toEnum 1 = InverseQuery
    toEnum 2 = ServerStatusRequest
    toEnum _ = undefined

instance Enum ResponseCode where
    fromEnum NoError        = 0
    fromEnum FormatError    = 1
    fromEnum ServerFailure  = 2
    fromEnum NameError      = 3
    fromEnum NotImplemented = 4
    fromEnum Refused        = 5

    toEnum 0 = NoError
    toEnum 1 = FormatError
    toEnum 2 = ServerFailure
    toEnum 3 = NameError
    toEnum 4 = NotImplemented
    toEnum 5 = Refused
    toEnum _ = undefined

instance Enum RecordType where
    fromEnum A       = 1
    fromEnum NS      = 2
    fromEnum MD      = 3
    fromEnum MF      = 4
    fromEnum CNAME   = 5
    fromEnum SOA     = 6
    fromEnum MB      = 7
    fromEnum MG      = 8
    fromEnum MR      = 9
    fromEnum NULL    = 10
    fromEnum WKS     = 11
    fromEnum PTR     = 12
    fromEnum HINFO   = 13
    fromEnum MINFO   = 14
    fromEnum MX      = 15
    fromEnum TXT     = 16
    fromEnum AXFR    = 252
    fromEnum MAILB   = 253
    fromEnum MAILA   = 254
    fromEnum AnyType = 255

    toEnum 1  = A
    toEnum 2  = NS
    toEnum 3  = MD
    toEnum 4  = MF
    toEnum 5  = CNAME
    toEnum 6  = SOA
    toEnum 7  = MB
    toEnum 8  = MG
    toEnum 9  = MR
    toEnum 10 = NULL
    toEnum 11 = WKS
    toEnum 12 = PTR
    toEnum 13 = HINFO
    toEnum 14 = MINFO
    toEnum 15 = MX
    toEnum 16 = TXT
    toEnum 252 = AXFR
    toEnum 253 = MAILB
    toEnum 254 = MAILA
    toEnum 255 = AnyType
    toEnum _  = undefined

instance Enum RecordClass
    fromEnum IN       = 1
    fromEnum CS       = 2
    fromEnum CH       = 3
    fromEnum HS       = 4
    fromEnum AnyClass = 255

    toEnum 1   = IN
    toEnum 2   = CS
    toEnum 3   = CH
    toEnum 4   = HS
    toEnum 255 = AnyClass
    toEnum _   = undefined