module Network.DNS.DomainMap
    ( DomainMap

    , empty
    , null
    , singleton
    , size

    , fromList
    , toList
    , toListBy

    , nearest
    , lookup

    , insert
    , delete
    )
    where

import qualified Data.List as L hiding (null)
import           Data.Maybe
import qualified Data.Trie as T
import           Data.Trie (Trie)
import           Network.DNS.Message
import           Prelude hiding (null, lookup)


newtype DomainMap a = DM (Node a)

data Node a = Node (Maybe a) (Trie (Node a))


empty :: DomainMap a
empty = DM emptyNode


emptyNode :: Node a
emptyNode = Node Nothing T.empty


null :: DomainMap a -> Bool
null (DM n) = nodeNull n


nodeNull :: Node a -> Bool
nodeNull (Node (Just _) _) = False
nodeNull (Node Nothing  t) = all nodeNull $ T.toListBy (\ _ a -> a) t


singleton :: DomainName -> a -> DomainMap a
singleton key value
    = insert key value empty


size :: forall a. DomainMap a -> Int
size (DM n) = nodeSize n


nodeSize :: Node a -> Int
nodeSize (Node m t)
    = let x  = if isJust m then
                   1
               else
                   0
          xs = map nodeSize $ T.toListBy (\ _ a -> a) t
      in
        L.foldl' (+) x xs


fromList :: forall a. [(DomainName, a)] -> DomainMap a
fromList []            = empty
fromList ((k, v) : xs) = insert k v $ fromList xs


toList :: forall a. DomainMap a -> [(DomainName, a)]
toList (DM root) = toList' rootName root
    where
      toList' :: DomainName -> Node a -> [(DomainName, a)]
      toList' soFar (Node m t)
          = let x  = case m of
                       Just v  -> [(soFar, v)]
                       Nothing -> []
                xs = concat $ map toList'' $ T.toList t
            in
              x ++ xs
          where
            toList'' :: (DomainLabel, Node a) -> [(DomainName, a)]
            toList'' (l, n) = toList' (consLabel l soFar) n


toListBy :: forall a b. DomainMap a -> (DomainName -> a -> b) -> [b]
toListBy dm f = map (uncurry f) $ toList dm


nearest :: forall a. DomainName -> DomainMap a -> Maybe a
nearest key (DM root) = nearest' key' root
    where
      key' :: [DomainLabel]
      key' = reverse $ nameToLabels key

      nearest' :: [DomainLabel] -> Node a -> Maybe a
      nearest' []     _          = error ("Illegal key: " ++ show key)
      nearest' (_:[]) (Node m _) = m
      nearest' (x:xs) (Node m t) = let nearer = do subNode' <- T.lookup x t
                                                   nearest' xs subNode'
                                   in
                                     if isJust nearer then
                                         nearer
                                     else
                                         m


lookup :: forall a. DomainName -> DomainMap a -> Maybe a
lookup key (DM root) = lookup' key' root
    where
      key' :: [DomainLabel]
      key' = reverse $ nameToLabels key

      lookup' :: [DomainLabel] -> Node a -> Maybe a
      lookup' []     _          = error ("Illegal key: " ++ show key)
      lookup' (_:[]) (Node m _) = m
      lookup' (x:xs) (Node _ t) = do subNode' <- T.lookup x t
                                     lookup' xs subNode'


insert :: forall a. DomainName -> a -> DomainMap a -> DomainMap a
insert key value (DM root) = DM $ insert' key' root
    where
      key' :: [DomainLabel]
      key' = reverse $ nameToLabels key

      insert' :: [DomainLabel] -> Node a -> Node a
      insert' []     _          = error ("Illegal key: " ++ show key)
      insert' (_:[]) (Node _ t) = Node (Just value) t
      insert' (x:xs) (Node m t) = let subNode' = case T.lookup x t of
                                                   Just subNode -> insert' xs subNode
                                                   Nothing      -> insert' xs emptyNode
                                      trie'    = T.insert x subNode' t
                                  in
                                    Node m trie'

delete :: forall a. DomainName -> DomainMap a -> DomainMap a
delete key (DM root) = DM $ delete' key' root
    where
      key' :: [DomainLabel]
      key' = reverse $ nameToLabels key

      delete' :: [DomainLabel] -> Node a -> Node a
      delete' []     _          = error ("Illegal key: " ++ show key)
      delete' (_:[]) (Node _ t) = Node Nothing t
      delete' (x:xs) (Node m t) = let subNode' = case T.lookup x t of
                                                   Just subNode -> delete' xs subNode
                                                   Nothing      -> emptyNode
                                      trie'    = if nodeNull subNode' then
                                                     T.delete x t
                                                 else
                                                     T.insert x subNode' t
                                  in
                                    Node m trie'