{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE RecordWildCards #-}

module Hans.IP4.Dhcp.Client (
    DhcpConfig(..),
    defaultDhcpConfig,
    DhcpLease(..),
    dhcpClient,
  ) where

import Hans.Device.Types (Device(devMac))
import Hans.IP4.Dhcp.Codec (SubnetMask(..))
import Hans.IP4.Dhcp.Packet
import Hans.IP4.Dhcp.Options
import Hans.IP4.Packet (IP4,pattern WildcardIP4,pattern BroadcastIP4,IP4Mask(..))
import Hans.IP4.RoutingTable(Route(..),RouteType(..))
import Hans.Lens
import Hans.Socket
           (UdpSocket,newUdpSocket,sClose,sendto,recvfrom,SockPort
           ,defaultSocketConfig)
import Hans.Serialize (runPutPacket)
import Hans.Threads (forkNamed)
import Hans.Time (toUSeconds)
import Hans.Types (NetworkStack,networkStack,addRoute,addNameServer4)

import           Control.Concurrent (threadDelay,killThread)
import           Control.Monad (guard)
import qualified Data.ByteString.Lazy as L
import           Data.Maybe (fromMaybe,mapMaybe)
import           Data.Serialize.Get (runGetLazy)
import           Data.Time.Clock (NominalDiffTime)
import           System.Random (randomIO,randomRIO)
import           System.Timeout (timeout)


-- | BOOTP server port.
bootps :: SockPort
bootps  = 67

-- | BOOTP client port.
bootpc :: SockPort
bootpc  = 68

mkXid :: IO Xid
mkXid  = do w <- randomIO
            return (Xid w)


renderMessage :: Dhcp4Message -> L.ByteString
renderMessage msg = runPutPacket 236 256 L.empty (putDhcp4Message msg)

data DhcpConfig = DhcpConfig { dcInitialTimeout :: !NominalDiffTime
                               -- ^ Initial timeout

                             , dcRetries :: !Int
                               -- ^ Number of retries

                             , dcDefaultRoute :: Bool
                               -- ^ Whether or not routing information received
                               -- from the DHCP server should be used as the
                               -- default route for the network stack.

                             , dcAutoRenew :: Bool
                               -- ^ Whether or not to fork a renew thread once
                               -- configuration information has been received.
                             }

defaultDhcpConfig :: DhcpConfig
defaultDhcpConfig  = DhcpConfig { dcInitialTimeout = 4.0
                                , dcRetries        = 6
                                , dcDefaultRoute   = True
                                , dcAutoRenew      = True }


-- | Wait for a result on a socket
waitResponse :: DhcpConfig -> IO () -> IO a -> IO (Maybe a)
waitResponse DhcpConfig { .. } send recv =
  go dcRetries (toUSeconds dcInitialTimeout)
  where
  go retries toVal =
    do send
       mb <- timeout toVal recv
       case mb of

         Just{} -> return mb

         -- adjust the timeout by two, and add some slack before trying again
         Nothing | retries > 0 ->
           do slack <- randomRIO (500,1000)
              go (retries - 1) (toVal * 2 + slack * 1000)

         _ -> return Nothing


data DhcpLease = DhcpLease { dhcpRenew :: !(IO ())
                           , dhcpAddr  :: !IP4
                           }


dhcpClient :: NetworkStack -> DhcpConfig -> Device -> IO (Maybe DhcpLease)
dhcpClient ns cfg dev =
  do sock <- newUdpSocket ns defaultSocketConfig (Just dev) WildcardIP4 (Just bootpc)
     dhcpDiscover cfg dev sock


-- | Discover a dhcp server, and request an address.
dhcpDiscover :: DhcpConfig -> Device -> UdpSocket IP4 -> IO (Maybe DhcpLease)
dhcpDiscover cfg dev sock =
  do xid  <- mkXid
     let msg = renderMessage (discoverToMessage (mkDiscover xid (devMac dev)))

     mb <- waitResponse cfg (sendto sock BroadcastIP4 bootps msg) (awaitOffer sock)
     case mb of
       Just offer -> dhcpRequest cfg dev sock offer
       Nothing    -> do sClose sock
                        return Nothing


-- | Only accept an offer.
awaitOffer :: UdpSocket IP4 -> IO Offer
awaitOffer sock = go
  where
  go =
    do (_,_,srcPort,bytes) <- recvfrom sock

       if srcPort /= bootps
          then go
          else case runGetLazy getDhcp4Message bytes of

                 Right msg
                   | Just (Right (OfferMessage o)) <- parseDhcpMessage msg ->
                     return o

                 _ -> go


-- | Respond to an offer with a request, and configure the network stack if an
-- acknowledgement is received.
dhcpRequest :: DhcpConfig -> Device -> UdpSocket IP4 -> Offer -> IO (Maybe DhcpLease)
dhcpRequest cfg dev sock offer =
  do let req = renderMessage (requestToMessage (offerToRequest offer))
     mb <- waitResponse cfg (sendto sock BroadcastIP4 bootps req) (awaitAck sock)
     sClose sock
     case mb of

       Nothing  -> return Nothing

       Just ack ->
         do lease <- handleAck (view networkStack sock) cfg dev offer ack
            return (Just lease)


awaitAck :: UdpSocket IP4 -> IO Ack
awaitAck sock = go
  where
  go =
    do (_,_,srcPort,bytes) <- recvfrom sock

       if srcPort /= bootps
          then go
          else case runGetLazy getDhcp4Message bytes of

                 Right msg
                   | Just (Right (AckMessage a)) <- parseDhcpMessage msg ->
                     return a

                 _ -> go



-- | Perform a DHCP Renew.
renew :: NetworkStack -> DhcpConfig -> Device -> Offer -> IO ()
renew ns cfg dev offer =
  do sock <- newUdpSocket ns defaultSocketConfig (Just dev) WildcardIP4 (Just bootpc)
     _    <- dhcpRequest cfg dev sock offer

     return ()


-- Ack Management --------------------------------------------------------------

-- | Apply the information in the Ack to the NetworkStack, and Device. Returns
-- information about the lease, as well as an IO action that can be used to
-- renew it.
handleAck :: NetworkStack -> DhcpConfig -> Device -> Offer -> Ack -> IO DhcpLease
handleAck ns cfg dev offer Ack { .. } =
  do let addr     = ackYourAddr
         mask     = fromMaybe 24 (lookupSubnet ackOptions)

     let nameServers = concat (mapMaybe getNameServers ackOptions)
     mapM_ (addNameServer4 ns) nameServers

     addRoute ns False Route
       { routeNetwork = IP4Mask addr mask
       , routeType    = Direct
       , routeDevice  = dev
       }

     case lookupGateway ackOptions of
       Just gw | dcDefaultRoute cfg ->
         addRoute ns True Route
             { routeNetwork = IP4Mask addr 0
             , routeType    = Indirect gw
             , routeDevice  = dev
             }

       _ -> return ()

     dhcpRenew <-
       if dcAutoRenew cfg
          then -- wait for half of the lease time, then automatically renew
               -- XXX: what happens on a 32-bit system here?
               do tid <- forkNamed "dhcpRenew" $
                         do threadDelay (fromIntegral ackLeaseTime * 500000)
                            renew ns cfg dev offer

                  return $ do killThread tid
                              renew ns cfg dev offer

          else return (renew ns cfg dev offer)

     return $! DhcpLease { dhcpAddr = addr, .. }


lookupGateway :: [Dhcp4Option] -> Maybe IP4
lookupGateway  = foldr p Nothing
  where
  p (OptRouters rs) _ = guard (not (null rs)) >> Just (head rs)
  p _               a = a

lookupSubnet :: [Dhcp4Option] -> Maybe Int
lookupSubnet  = foldr p Nothing
  where
  p (OptSubnetMask (SubnetMask i)) _ = Just i
  p _                              a = a

getNameServers :: Dhcp4Option -> Maybe [IP4]
getNameServers (OptNameServers addrs) = Just addrs
getNameServers _                      = Nothing