module Network.DNS.Message
    ( Message(..)
    , MessageID
    , MessageType(..)
    , Header(..)
    , Opcode(..)
    , ResponseCode(..)
    , Question(..)
    , ResourceRecord(..)
    , DomainName
    , DomainLabel
    , TTL
    , SomeRR(..)
    , RecordType(..)
    , RecordClass(..)

    , CNAME(..)
    , HINFO(..)
    )
    where

import           Control.Monad
import           Data.Binary
import           Data.Binary.BitPut as BP
import           Data.Binary.Get as G
import           Data.Binary.Put as P
import           Data.Binary.Strict.BitGet as BG
import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy as LBS
import           Data.Typeable
import qualified Data.IntMap as IM
import           Data.IntMap (IntMap)
import           Data.Word


replicateM' :: Monad m => Int -> (a -> m (b, a)) -> a -> m ([b], a)
replicateM' = worker []
    where
      worker :: Monad m => [b] -> Int -> (a -> m (b, a)) -> a -> m ([b], a)
      worker soFar 0 _ a = return (reverse soFar, a)
      worker soFar n f a = do (b, a') <- f a
                              worker (b : soFar) (n - 1) f a'


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

data Header
    = Header {
        hdMessageID             :: !MessageID
      , 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
      }

type MessageID = Word16

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)

putQ :: Question -> Put
putQ q
    = do putDomainName $ qName q
         put $ qType  q
         put $ qClass q

getQ :: DecompTable -> Get (Question, DecompTable)
getQ dt
    = do (nm, dt') <- getDomainName dt
         ty        <- get
         cl        <- get
         let q = Question {
                   qName  = nm
                 , qType  = ty
                 , qClass = cl
                 }
         return (q, dt')

type DomainName  = [DomainLabel]
type DomainLabel = BS.ByteString

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

class (Typeable rr, Show rr, Eq rr) => ResourceRecord rr where
    rrName    :: rr -> DomainName
    rrType    :: rr -> RecordType
    rrClass   :: rr -> RecordClass
    rrTTL     :: rr -> TTL
    rrPutData :: rr -> Put
    rrGetData :: DecompTable -> DomainName -> RecordClass -> TTL -> Get (rr, DecompTable)
    toRR      :: rr -> SomeRR
    fromRR    :: SomeRR -> Maybe rr

    toRR   rr           = SomeRR rr
    fromRR (SomeRR rr') = cast rr'

putRR :: ResourceRecord rr => rr -> Put
putRR rr = do putDomainName $ rrName rr
              put $ rrType  rr
              put $ rrClass rr
              putWord32be $ rrTTL rr

              let dat = runPut $ rrPutData rr
              putWord16be $ fromIntegral $ LBS.length dat
              putLazyByteString dat

getRR :: DecompTable -> Get (SomeRR, DecompTable)
getRR dt
    = do (nm, dt') <- getDomainName dt
         ty        <- get
         cl        <- get
         ttl       <- G.getWord32be
         case ty of
           CNAME   -> do (rr, dt'') <- rrGetData dt' nm cl ttl
                         return (toRR (rr :: CNAME), dt'')
           HINFO   -> do (rr, dt'') <- rrGetData dt' nm cl ttl
                         return (toRR (rr :: HINFO), dt'')
           AXFR    -> onlyForQuestions "AXFR"
           MAILB   -> onlyForQuestions "MAILB"
           MAILA   -> onlyForQuestions "MAILA"
           AnyType -> onlyForQuestions "ANY"
    where
      onlyForQuestions name
          = fail (name ++ " is only for questions, not an actual resource record.")

data SomeRR = forall rr. ResourceRecord rr => SomeRR rr
              deriving Typeable
instance ResourceRecord SomeRR where
    rrName    (SomeRR rr) = rrName  rr
    rrType    (SomeRR rr) = rrType  rr
    rrClass   (SomeRR rr) = rrClass rr
    rrTTL     (SomeRR rr) = rrTTL   rr
    rrPutData (SomeRR rr) = rrPutData rr
    rrGetData _ _ _ _     = fail "SomeRR can't directly be constructed."
    toRR   = id
    fromRR = Just
instance Eq SomeRR where
    (SomeRR a) == (SomeRR b) = Just a == cast b
instance Show SomeRR where
    show (SomeRR rr) = show rr

type DecompTable = IntMap BS.ByteString
type TTL = Word32

data CNAME = CNAME' !DomainName !RecordClass !TTL !DomainName
             deriving (Eq, Show, Typeable)
instance ResourceRecord CNAME where
    rrName    (CNAME' n _ _ _) = n
    rrType    _                = CNAME
    rrClass   (CNAME' _ c _ _) = c
    rrTTL     (CNAME' _ _ t _) = t
    rrGetData dt n c t         = do (d, dt') <- getDomainName dt
                                    return (CNAME' n c t d, dt')
    rrPutData (CNAME' _ _ _ d) = putDomainName d

data HINFO = HINFO' !DomainName !RecordClass !TTL !BS.ByteString !BS.ByteString
             deriving (Eq, Show, Typeable)
instance ResourceRecord HINFO where
    rrName    (HINFO' n _ _ _ _) = n
    rrType    _                  = HINFO
    rrClass   (HINFO' _ c _ _ _) = c
    rrTTL     (HINFO' _ _ t _ _) = t
    rrGetData dt n c t           = do cpu <- getCharString
                                      os  <- getCharString
                                      return (HINFO' n c t cpu os, dt)
    rrPutData (HINFO' _ _ _ c o) = do putCharString c
                                      putCharString o

getDomainName :: DecompTable -> Get (DomainName, DecompTable)
getDomainName = flip worker []
    where
      worker :: DecompTable -> [DomainLabel] -> Get ([DomainLabel], DecompTable)
      worker dt soFar
          = do (l, dt') <- getDomainLabel dt
               case BS.null l of
                 True  -> return (reverse (l : soFar), dt')
                 False -> worker dt' (l : soFar)

getDomainLabel :: DecompTable -> Get (DomainLabel, DecompTable)
getDomainLabel dt
    = do header <- getByteString 1
         let Right h
                 = runBitGet header $
                   do a <- getBit
                      b <- getBit
                      n <- liftM fromIntegral (getAsWord8 6)
                      case (a, b) of
                        ( True,  True) -> return $ Offset n
                        (False, False) -> return $ Length n
                        _              -> fail "Illegal label header"
         case h of
           Offset n
               -> do let Just l = IM.lookup n dt
                     return (l, dt)
           Length n
               -> do offset <- liftM fromIntegral bytesRead
                     label  <- getByteString n
                     let dt' = IM.insert offset label dt
                     return (label, dt')

getCharString :: Get BS.ByteString
getCharString = do len <- G.getWord8
                   getByteString (fromIntegral len)

putCharString :: BS.ByteString -> Put
putCharString = putDomainLabel

data LabelHeader
    = Offset !Int
    | Length !Int

putDomainName :: DomainName -> Put
putDomainName = mapM_ putDomainLabel

putDomainLabel :: DomainLabel -> Put
putDomainLabel l
    = do putWord8 $ fromIntegral $ BS.length l
         P.putByteString l

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 Message where
    put m = do put $ msgHeader m
               putWord16be $ fromIntegral $ length $ msgQuestions m
               putWord16be $ fromIntegral $ length $ msgAnswers m
               putWord16be $ fromIntegral $ length $ msgAuthorities m
               putWord16be $ fromIntegral $ length $ msgAdditionals m
               mapM_ putQ  $ msgQuestions m
               mapM_ putRR $ msgAnswers m
               mapM_ putRR $ msgAuthorities m
               mapM_ putRR $ msgAdditionals m

    get = do hdr  <- get
             nQ   <- liftM fromIntegral G.getWord16be
             nAns <- liftM fromIntegral G.getWord16be
             nAth <- liftM fromIntegral G.getWord16be
             nAdd <- liftM fromIntegral G.getWord16be
             (qs  , dt1) <- replicateM' nQ   getQ  IM.empty
             (anss, dt2) <- replicateM' nAns getRR dt1
             (aths, dt3) <- replicateM' nAth getRR dt2
             (adds, _  ) <- replicateM' nAdd getRR dt3
             return Message {
                          msgHeader      = hdr
                        , msgQuestions   = qs
                        , msgAnswers     = anss
                        , msgAuthorities = aths
                        , msgAdditionals = adds
                        }

instance Binary Header where
    put h = do putWord16be $ hdMessageID h
               putLazyByteString flags
        where
          flags = runBitPut $
                  do putNBits 1 $ fromEnum $ hdMessageType h
                     putNBits 4 $ fromEnum $ hdOpcode h
                     putBit $ hdIsAuthoritativeAnswer h
                     putBit $ hdIsTruncated h
                     putBit $ hdIsRecursionDesired h
                     putBit $ hdIsRecursionAvailable h
                     putNBits 3 (0 :: Int)
                     putNBits 4 $ fromEnum $ hdResponseCode h

    get = do mID   <- G.getWord16be
             flags <- getByteString 2
             let Right hd
                     = runBitGet flags $
                       do qr <- liftM (toEnum . fromIntegral) $ getAsWord8 1
                          op <- liftM (toEnum . fromIntegral) $ getAsWord8 4
                          aa <- getBit
                          tc <- getBit
                          rd <- getBit
                          ra <- getBit
                          BG.skip 3
                          rc <- liftM (toEnum . fromIntegral) $ getAsWord8 4
                          return 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 where
    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

instance Binary RecordType where
    get = liftM (toEnum . fromIntegral) G.getWord16be
    put = putWord16be . fromIntegral . fromEnum

instance Binary RecordClass where
    get = liftM (toEnum . fromIntegral) G.getWord16be
    put = putWord16be . fromIntegral . fromEnum