    {-# LANGUAGE CPP           #-}
    {-# LANGUAGE DeriveGeneric #-}

    {-|
      IP routing table is a tree of 'AddrRange'
      to search one of them on the longest
      match base. It is a kind of TRIE with one
      way branching removed. Both IPv4 and IPv6
      are supported.
    -}
    module Data.IP.RouteTable.Internal where

    import Control.Applicative hiding (empty)
    import qualified Control.Applicative as A (empty)
    import Control.Monad
    import Data.Bits
    import Data.Foldable (Foldable(..))
    import Data.IP.Addr
    import Data.IP.Op
    import Data.IP.Range
    import Data.IntMap (IntMap, (!))
    import qualified Data.IntMap as IM (fromList)
    import Data.Monoid hiding ((<>))
    import Data.Semigroup
    import Data.Traversable
    import Data.Word
    import GHC.Generics (Generic, Generic1)
    import Prelude hiding (lookup)

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

    {-|
      A class to contain IPv4 and IPv6.
    -}
    class Addr a => Routable a where
        {-|
          The 'intToTBit' function takes 'Int' and returns an 'Routable' address
          whose only n-th bit is set.
        -}
        intToTBit   :: Int -> a
        {-|
          The 'isZero' function takes an 'Routable' address and an test bit
          'Routable' address and returns 'True' is the bit is unset,
          otherwise returns 'False'.
        -}
        isZero :: a -> a -> Bool

    instance Routable IPv4 where
        intToTBit = intToTBitIPv4
        isZero a b = a `masked` b == IP4 0

    instance Routable IPv6 where
        intToTBit = intToTBitIPv6
        isZero a b = a `masked` b == IP6 (0,0,0,0)

    ----------------------------------------------------------------
    --
    -- Test Bit
    --

    intToTBitIPv4 :: Int -> IPv4
    intToTBitIPv4 len = IP4 (intToTBitsIPv4 ! len)

    intToTBitIPv6 :: Int -> IPv6
    intToTBitIPv6 len = IP6 (intToTBitsIPv6 ! len)

    intToTBitsWord32 :: [Word32]
    intToTBitsWord32 = iterate (`shift` (-1)) 0x80000000

    intToTBitsIPv4 :: IntMap IPv4Addr
    intToTBitsIPv4 = IM.fromList $ zip [0..32] intToTBitsWord32

    intToTBitsIPv6 :: IntMap IPv6Addr
    intToTBitsIPv6 = IM.fromList $ zip [0..128] bs
      where
        bs = b1 ++ b2 ++ b3 ++ b4 ++ b5
        b1 = map (\vbit -> (vbit,all0,all0,all0)) intToTBits
        b2 = map (\vbit -> (all0,vbit,all0,all0)) intToTBits
        b3 = map (\vbit -> (all0,all0,vbit,all0)) intToTBits
        b4 = map (\vbit -> (all0,all0,all0,vbit)) intToTBits
        b5 =              [(all0,all0,all0,all0)]
        intToTBits = take 32 intToTBitsWord32
        all0 = 0x00000000

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

    {-|
      The Tree structure for IP routing table based on TRIE with
      one way branching removed. This is an abstract data type,
      so you cannot touch its inside. Please use 'insert' or 'lookup', instead.
    -}
    data IPRTable k a =
        Nil
      | Node !(AddrRange k) !k !(Maybe a) !(IPRTable k a) !(IPRTable k a)
      deriving (Eq, Generic, Generic1, Show)

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

    {-|
      The 'empty' function returns an empty IP routing table.

    >>> (empty :: IPRTable IPv4 ()) == fromList []
    True
    -}
    empty :: Routable k => IPRTable k a
    empty = Nil

    instance Functor (IPRTable k) where
        fmap _ Nil = Nil
        fmap f (Node r a mv b1 b2) = Node r a (f <$> mv) (fmap f b1) (fmap f b2)

    instance Foldable (IPRTable k) where
        foldMap _ Nil = mempty
        foldMap f (Node _ _ mv b1 b2) = foldMap f mv `mappend` foldMap f b1 `mappend` foldMap f b2

    instance Traversable (IPRTable k) where
        traverse _ Nil = pure Nil
        traverse f (Node r a mv b1 b2) = Node r a <$> traverse f mv <*> traverse f b1 <*> traverse f b2

    -- | Note that Semigroup and Monoid instances are right-biased.
    --   That is, if both arguments have the same key, the value from the right
    --   argument will be used.
    --   Since: 1.7.5
    instance Routable k => Semigroup (IPRTable k a) where
        a <> b = foldlWithKey (\rt k v -> insert k v rt) a b
        stimes = stimesIdempotent

    -- | Since: 1.7.5
    instance Routable k => Monoid (IPRTable k a) where
        mempty = empty
        mappend = (<>)

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

    {-|
      The 'insert' function inserts a value with a key of 'AddrRange' to 'IPRTable'
      and returns a new 'IPRTable'.

    >>> (insert ("127.0.0.1" :: AddrRange IPv4) () empty) == fromList [("127.0.0.1",())]
    True
    -}
    insert :: (Routable k) => AddrRange k -> a -> IPRTable k a -> IPRTable k a
    insert k1 v1 Nil = Node k1 tb1 (Just v1) Nil Nil
      where
        tb1 = keyToTestBit k1
    insert k1 v1 s@(Node k2 tb2 v2 l r)
      | k1 == k2  = Node k1 tb1 (Just v1) l r
      | k2 >:> k1 = if isLeft k1 tb2 then
                        Node k2 tb2 v2 (insert k1 v1 l) r
                      else
                        Node k2 tb2 v2 l (insert k1 v1 r)
      | k1 >:> k2 = if isLeft k2 tb1 then
                        Node k1 tb1 (Just v1) s Nil
                      else
                        Node k1 tb1 (Just v1) Nil s
      | otherwise = let n = Node k1 tb1 (Just v1) Nil Nil
                    in link n s
      where
        tb1 = keyToTestBit k1

    link :: Routable k => IPRTable k a -> IPRTable k a -> IPRTable k a
    link s1@(Node k1 _ _ _ _) s2@(Node k2 _ _ _ _)
      | isLeft k1 tbg = Node kg tbg Nothing s1 s2
      | otherwise     = Node kg tbg Nothing s2 s1
      where
        kg = glue 0 k1 k2
        tbg = keyToTestBit kg
    link _ _ = error "link"

    glue :: (Routable k) => Int -> AddrRange k -> AddrRange k -> AddrRange k
    glue n k1 k2
      | addr k1 `masked` mk == addr k2 `masked` mk = glue (n + 1) k1 k2
      | otherwise = makeAddrRange (addr k1) (n - 1)
      where
        mk = intToMask n

    keyToTestBit :: Routable k => AddrRange k -> k
    keyToTestBit = intToTBit . mlen

    isLeft :: Routable k => AddrRange k -> k -> Bool
    isLeft adr = isZero (addr adr)

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

    {-|
      The 'delete' function deletes a value by a key of 'AddrRange' from 'IPRTable'
      and returns a new 'IPRTable'.

    >>> delete "127.0.0.1" (insert "127.0.0.1" () empty) == (empty :: IPRTable IPv4 ())
    True
    -}
    delete :: (Routable k) => AddrRange k -> IPRTable k a -> IPRTable k a
    delete _ Nil = Nil
    delete k1 s@(Node k2 tb2 v2 l r)
      | k1 == k2  = node k2 tb2 Nothing l r
      | k2 >:> k1 = if isLeft k1 tb2 then
                        node k2 tb2 v2 (delete k1 l) r
                      else
                        node k2 tb2 v2 l (delete k1 r)
      | otherwise = s

    node :: (Routable k) => AddrRange k -> k -> Maybe a -> IPRTable k a -> IPRTable k a -> IPRTable k a
    node _ _ Nothing Nil r = r
    node _ _ Nothing l Nil = l
    node k tb v      l   r = Node k tb v l r

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

    {-|
      The 'lookup' function looks up 'IPRTable' with a key of 'AddrRange'.
      If a routing information in 'IPRTable' matches the key, its value
      is returned.

    >>> let v4 = ["133.4.0.0/16","133.5.0.0/16","133.5.16.0/24","133.5.23.0/24"] :: [AddrRange IPv4]
    >>> let rt = fromList $ zip v4 v4
    >>> lookup "127.0.0.1" rt
    Nothing
    >>> lookup "133.3.0.1" rt
    Nothing
    >>> lookup "133.4.0.0" rt
    Just 133.4.0.0/16
    >>> lookup "133.4.0.1" rt
    Just 133.4.0.0/16
    >>> lookup "133.5.16.0" rt
    Just 133.5.16.0/24
    >>> lookup "133.5.16.1" rt
    Just 133.5.16.0/24
    -}
    lookup :: Routable k => AddrRange k -> IPRTable k a -> Maybe a
    lookup k s = fmap snd (search k s Nothing)

    {-|
      The 'lookupKeyValue' function looks up 'IPRTable' with a key of 'AddrRange'.
      If a routing information in 'IPRTable' matches the key, both key and value
      are returned.

    >>> :set -XOverloadedStrings
    >>> let rt = fromList ([("192.168.0.0/24", 1), ("10.10.0.0/16", 2)] :: [(AddrRange IPv4, Int)])
    >>> lookupKeyValue "127.0.0.1" rt
    Nothing
    >>> lookupKeyValue "192.168.0.1" rt
    Just (192.168.0.0/24,1)
    >>> lookupKeyValue "10.10.0.1" rt
    Just (10.10.0.0/16,2)
    -}
    lookupKeyValue :: Routable k => AddrRange k -> IPRTable k a -> Maybe (AddrRange k, a)
    lookupKeyValue k s = search k s Nothing

    search :: Routable k => AddrRange k
                         -> IPRTable k a
                         -> Maybe (AddrRange k, a)
                         -> Maybe (AddrRange k, a)
    search _ Nil res = res
    search k1 (Node k2 tb2 Nothing l r) res
      | k1 == k2  = res
      | k2 >:> k1 = if isLeft k1 tb2 then
                        search k1 l res
                      else
                        search k1 r res
      | otherwise = res
    search k1 (Node k2 tb2 (Just vl) l r) res
      | k1 == k2  = Just (k1, vl)
      | k2 >:> k1 = if isLeft k1 tb2 then
                        search k1 l $ Just (k2, vl)
                      else
                        search k1 r $ Just (k2, vl)
      | otherwise = res

    {-|
      'lookupAll' is a version of 'lookup' that returns all entries matching the
       given key, not just the longest match.

    >>> :set -XOverloadedStrings
    >>> let rt = fromList ([("192.168.0.0/24", 1), ("10.10.0.0/16", 2), ("10.0.0.0/8", 3)] :: [(AddrRange IPv4, Int)])
    >>> lookupAll "127.0.0.1" rt
    []
    >>> lookupAll "192.168.0.1" rt
    [(192.168.0.0/24,1)]
    >>> lookupAll "10.10.0.1" rt
    [(10.10.0.0/16,2),(10.0.0.0/8,3)]
    -}

    lookupAll :: Routable k => AddrRange k -> IPRTable k a -> [(AddrRange k, a)]
    lookupAll range = go []
      where
        go acc Nil = acc
        go acc (Node k tb Nothing l r)
          | k == range = acc
          | k >:> range = go acc $ if isLeft range tb then l else r
          | otherwise = acc
        go acc (Node k tb (Just v) l r)
          | k == range = (k,v):acc
          | k >:> range = go ((k,v):acc) $ if isLeft range tb then l else r
          | otherwise = acc

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

    {-|
      The 'findMatch' function looks up 'IPRTable' with a key of 'AddrRange'.
      If the key matches routing informations in 'IPRTable', they are
      returned.

    >>> let v4 = ["133.4.0.0/16","133.5.0.0/16","133.5.16.0/24","133.5.23.0/24"] :: [AddrRange IPv4]
    >>> let rt = fromList $ zip v4 $ repeat ()
    >>> findMatch "133.4.0.0/15" rt :: [(AddrRange IPv4,())]
    [(133.4.0.0/16,()),(133.5.0.0/16,()),(133.5.16.0/24,()),(133.5.23.0/24,())]
    -}

    findMatch :: Alternative m => Routable k => AddrRange k -> IPRTable k a -> m (AddrRange k, a)
    findMatch _ Nil = A.empty
    findMatch k1 (Node k2 _ Nothing l r)
      | k1 >:> k2 = findMatch k1 l <|> findMatch k1 r
      | k2 >:> k1 = findMatch k1 l <|> findMatch k1 r
      | otherwise = A.empty
    findMatch k1 (Node k2 _ (Just vl) l r)
      | k1 >:> k2 = pure (k2, vl) <|> findMatch k1 l <|> findMatch k1 r
      | k2 >:> k1 = findMatch k1 l <|> findMatch k1 r
      | otherwise = A.empty

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

    {-|
      The 'fromList' function creates a new IP routing table from
      a list of a pair of 'IPrange' and value.
    -}
    fromList :: Routable k => [(AddrRange k, a)] -> IPRTable k a
    fromList = foldl' (\s (k,v) -> insert k v s) empty

    {-|
      The 'toList' function creates a list of a pair of 'AddrRange' and
      value from an IP routing table.
    -}
    toList :: Routable k => IPRTable k a -> [(AddrRange k, a)]
    toList = foldt toL []
      where
        toL Nil xs = xs
        toL (Node _ _ Nothing  _ _) xs = xs
        toL (Node k _ (Just a) _ _) xs = (k,a) : xs

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

    foldt :: (IPRTable k a -> b -> b) -> b -> IPRTable k a -> b
    foldt _ v Nil = v
    foldt func v rt@(Node _ _ _ l r) = foldt func (foldt func (func rt v) l) r

    -- | /O(n)/. Fold the keys and values in the IPRTable using the given
    --   left-associative binary operator.
    --   This function is equivalent to Data.Map.foldlWithKey with necessary to
    --   IPRTable changes.
    --   Since: 1.7.5
    foldlWithKey :: (b -> AddrRange k -> a -> b) -> b -> IPRTable k a -> b
    foldlWithKey f zr = go zr
      where
        go z Nil = z
        go z (Node _ _ Nothing l r) = go (go z l) r
        go z (Node n _ (Just v) l r) = go (f (go z l) n v) r
    {-# INLINE foldlWithKey #-}

    -- | /O(n)/. Fold the keys and values in the IPRTable using the given
    --   right-associative binary operator.
    --   This function is equivalent to Data.Map.foldrWithKey with necessary to
    --   IPRTable changes.
    --   Since: 1.7.5
    foldrWithKey :: (AddrRange k -> a -> b -> b) -> b -> IPRTable k a -> b
    foldrWithKey f zr = go zr
      where
        go z Nil = z
        go z (Node _ _ Nothing l r) = go (go z r) l
        go z (Node n _ (Just v) l r) = go (f n v (go z r)) l
    {-# INLINE foldrWithKey #-}