module Hawl.Data.PathTree (
     Tree(..), 
     empty, fromList, insert,
     lookup, toList, addPrefix
    ) where

import Hawl.Data.Path

import Data.List (foldl')
import Data.Map (Map)
import qualified Data.Map as Map
import Prelude hiding (lookup)


data Tree a = Node (Map String (Tree a)) | Leaf a

-- | An empty tree.
empty :: Tree a
empty = Node Map.empty

singleton :: Path -> a -> Tree a
singleton ps x = insert ps x empty

-- | Construct a path tree from a list of (path, value) pairs.
--   Uses the last value if there a duplicate paths. Note that if a 
--   leaf L2 with a path p comes after a leaf L1 with a path which has p as a prefix,
--   L1 disappears.
fromList :: [(Path, a)] -> Tree a
fromList = foldl' (flip (uncurry insert)) empty

insert :: Path -> a -> Tree a -> Tree a
insert [] x _ = Leaf x
insert (p:ps) x (Node m) = Node (Map.insertWith u p (singleton ps x) m)
    where u = const (insert ps x)
insert ps x (Leaf _) = singleton ps x

lookup :: Monad m => Path -> Tree a -> m a
lookup p = lookup_ p
    where 
      lookup_ []     (Leaf x) = return x
      lookup_ (c:cs) (Node m) = maybe notFound (lookup_ cs) (Map.lookup c m)
      lookup_ _      _        = notFound
      notFound = fail $ "Not found: " ++ showPath p

toList :: Tree a -> [(Path,a)]
toList (Node m) = [ (c:p, x) | (c,t) <- Map.toList m, (p,x) <- toList t]
toList (Leaf x) = [([],x)]

-- | Add a prefix to all path in a tree.
addPrefix :: Path  -- ^ Prefix to add to the paths.
          -> Tree a -> Tree a
addPrefix = flip $ foldl' $ \t c -> Node (Map.singleton c t)