{-# LANGUAGE ScopedTypeVariables, MultiParamTypeClasses, FlexibleInstances #-}

-- | The abstract representation of a Tree and useful abstract utilities to
-- handle those.
module Storage.Hashed.Tree
    ( Tree, Blob(..), TreeItem(..), ItemType(..), Hash(..)
    , makeTree, makeTreeWithHash, emptyTree, emptyBlob, makeBlob, makeBlobBS

    -- * Unfolding stubbed (lazy) Trees.
    --
    -- | By default, Tree obtained by a read function is stubbed: it will
    -- contain Stub items that need to be executed in order to access the
    -- respective subtrees. 'expand' will produce an unstubbed Tree.
    , expandUpdate, expand, expandPath

    -- * Tree access and lookup.
    , items, list, listImmediate, treeHash
    , lookup, find, findFile, findTree, itemHash, itemType
    , zipCommonFiles, zipFiles, zipTrees, diffTrees

    -- * Files (Blobs).
    , readBlob

    -- * Filtering trees.
    , FilterTree(..), restrict

    -- * Manipulating trees.
    , modifyTree, updateTree, updateSubtrees, overlay ) where

import Prelude hiding( lookup, filter, all )
import Storage.Hashed.AnchoredPath
import Storage.Hashed.Hash

import qualified Data.ByteString.Lazy.Char8 as BL
import qualified Data.ByteString.Char8 as BS
import qualified Data.Map as M

import Data.Maybe( catMaybes )
import Data.List( union, sort )
import Control.Applicative( (<$>) )

--------------------------------
-- Tree, Blob and friends
--

data Blob m = Blob !(m BL.ByteString) !Hash
data TreeItem m = File !(Blob m)
                | SubTree !(Tree m)
                | Stub !(m (Tree m)) !Hash

data ItemType = BlobType | TreeType deriving (Show, Eq)

-- | Abstraction of a filesystem tree.
-- Please note that the Tree returned by the respective read operations will
-- have TreeStub items in it. To obtain a Tree without such stubs, call
-- expand on it, eg.:
--
-- > tree <- readDarcsPristine "." >>= expand
--
-- When a Tree is expanded, it becomes \"final\". All stubs are forced and the
-- Tree can be traversed purely. Access to actual file contents stays in IO
-- though.
--
-- A Tree may have a Hash associated with it. A pair of Tree's is identical
-- whenever their hashes are (the reverse need not hold, since not all Trees
-- come equipped with a hash).
data Tree m = Tree { items :: (M.Map Name (TreeItem m))
                   -- | Get hash of a Tree. This is guaranteed to uniquely
                   -- identify the Tree (including any blob content), as far as
                   -- cryptographic hashes are concerned. Sha256 is recommended.
                   , treeHash :: !Hash }

listImmediate :: Tree m -> [(Name, TreeItem m)]
listImmediate = M.toList . items

-- | Get a hash of a TreeItem. May be Nothing.
itemHash :: TreeItem m -> Hash
itemHash (File (Blob _ h)) = h
itemHash (SubTree t) = treeHash t
itemHash (Stub _ h) = h

itemType :: TreeItem m -> ItemType
itemType (File _) = BlobType
itemType (SubTree _) = TreeType
itemType (Stub _ _) = TreeType

emptyTree :: (Monad m) => Tree m
emptyTree = Tree { items = M.empty
                 , treeHash = NoHash }

emptyBlob :: (Monad m) => Blob m
emptyBlob = Blob (return BL.empty) NoHash

makeBlob :: (Monad m) => BL.ByteString -> Blob m
makeBlob str = Blob (return str) (sha256 str)

makeBlobBS :: (Monad m) => BS.ByteString -> Blob m
makeBlobBS s' = let s = BL.fromChunks [s'] in Blob (return s) (sha256 s)

makeTree :: (Monad m) => [(Name,TreeItem m)] -> Tree m
makeTree l = Tree { items = M.fromList l
                  , treeHash = NoHash }

makeTreeWithHash :: (Monad m) => [(Name,TreeItem m)] -> Hash -> Tree m
makeTreeWithHash l h = Tree { items = M.fromList l
                            , treeHash = h }

-----------------------------------
-- Tree access and lookup
--

-- | Look up a 'Tree' item (an immediate subtree or blob).
lookup :: Tree m -> Name -> Maybe (TreeItem m)
lookup t n = M.lookup n (items t)

find' :: TreeItem m -> AnchoredPath -> Maybe (TreeItem m)
find' t (AnchoredPath []) = Just t
find' (SubTree t) (AnchoredPath (d : rest)) =
    case lookup t d of
      Just sub -> find' sub (AnchoredPath rest)
      Nothing -> Nothing
find' _ _ = Nothing

-- | Find a 'TreeItem' by its path. Gives 'Nothing' if the path is invalid.
find :: Tree m -> AnchoredPath -> Maybe (TreeItem m)
find = find' . SubTree

-- | Find a 'Blob' by its path. Gives 'Nothing' if the path is invalid, or does
-- not point to a Blob.
findFile :: Tree m -> AnchoredPath -> Maybe (Blob m)
findFile t p = case find t p of
                 Just (File x) -> Just x
                 _ -> Nothing

-- | Find a 'Tree' by its path. Gives 'Nothing' if the path is invalid, or does
-- not point to a Tree.
findTree :: Tree m -> AnchoredPath -> Maybe (Tree m)
findTree t p = case find t p of
                 Just (SubTree x) -> Just x
                 _ -> Nothing

-- | List all contents of a 'Tree'.
list :: Tree m -> [(AnchoredPath, TreeItem m)]
list t_ = paths t_ (AnchoredPath [])
    where paths t p = [ (appendPath p n, i)
                          | (n,i) <- listImmediate t ] ++
                    concat [ paths subt (appendPath p subn)
                             | (subn, SubTree subt) <- listImmediate t ]

expandUpdate :: (Monad m) => (AnchoredPath -> Tree m -> m (Tree m)) -> Tree m -> m (Tree m)
expandUpdate update t_ = go (AnchoredPath []) t_
    where go path t = do
            let subtree (name, sub) = do tree <- go (path `appendPath` name) =<< unstub sub
                                         return (name, SubTree tree)
            expanded <- mapM subtree [ x | x@(_, item) <- listImmediate t, isSub item ]
            let orig = [ i | i <- listImmediate t, not $ isSub $ snd i ]
                orig_map = M.filter (not . isSub) (items t)
                expanded_map = M.fromList expanded
                tree = t { items = M.union orig_map expanded_map }
            update path tree
          unstub (Stub s _) = s
          unstub (SubTree t) = return t
          isSub (File _) = False
          isSub _ = True

-- | Expand a stubbed Tree into a one with no stubs in it. You might want to
-- filter the tree before expanding to save IO. This is the basic
-- implementation, which may be overriden by some Tree instances (this is
-- especially true of the Index case).
expand :: (Monad m) => Tree m -> m (Tree m)
expand = expandUpdate $ \_ -> return

-- | Unfold a path in a (stubbed) Tree, such that the leaf node of the path is
-- reachable without crossing any stubs.
expandPath :: (Monad m) => Tree m -> AnchoredPath -> m (Tree m)
expandPath t_ path_ = do expand' t_ path_
    where expand' t (AnchoredPath [_]) = return t
          expand' t (AnchoredPath (n:rest)) = do
            case lookup t n of
              (Just (Stub stub _)) ->
                  do unstubbed <- stub
                     amend t n rest unstubbed
              (Just (SubTree t')) -> amend t n rest t'
              _ -> fail $ "Descent error in expandPath: " ++ show path_
          amend t name rest sub = do
            sub' <- expand' sub (AnchoredPath rest)
            let orig_l = [ i | i@(n',_) <- listImmediate t, name /= n' ]
                tree = t { items = M.insert name (SubTree sub') (items t) }
            return tree

class (Monad m) => FilterTree a m where
    -- | Given @pred tree@, produce a 'Tree' that only has items for which
    -- @pred@ returns @True@.
    -- The tree might contain stubs. When expanded, these will be subject to
    -- filtering as well.
    filter :: (AnchoredPath -> TreeItem m -> Bool) -> a m -> a m

instance (Monad m) => FilterTree Tree m where
    filter predicate t_ = filter' t_ (AnchoredPath [])
        where filter' t path =
                  let subs = (catMaybes [ (,) name `fmap` wibble path name item
                                              | (name,item) <- listImmediate t ])
                  in t { items = M.mapMaybeWithKey (wibble path) $ items t }
              wibble path name item =
                  let npath = path `appendPath` name in
                      if predicate npath item
                         then Just $ filterSub npath item
                         else Nothing
              filterSub npath (SubTree t) = SubTree $ filter' t npath
              filterSub npath (Stub stub h) =
                  Stub (do x <- stub
                           return $ filter' x npath) h
              filterSub _ x = x

-- | Given two Trees, a @guide@ and a @tree@, produces a new Tree that is a
-- identical to @tree@, but only has those items that are present in both
-- @tree@ and @guide@. The @guide@ Tree may not contain any stubs.
restrict :: (FilterTree t m, Monad n) => Tree n -> t m -> t m
restrict guide tree = filter accept tree
    where accept path item =
              case (find guide path, item) of
                (Just (SubTree _), SubTree _) -> True
                (Just (SubTree _), Stub _ _) -> True
                (Just (File _), File _) -> True
                (Just (Stub _ _), _) ->
                    error "*sulk* Go away, you, you precondition violator!"
                (_, _) -> False

-- | Read a Blob into a Lazy ByteString. Might be backed by an mmap, use with
-- care.
readBlob :: Blob m -> m BL.ByteString
readBlob (Blob r _) = r

-- | For every pair of corresponding blobs from the two supplied trees,
-- evaluate the supplied function and accumulate the results in a list. Hint:
-- to get IO actions through, just use sequence on the resulting list.
-- NB. This won't expand any stubs.
zipCommonFiles :: (AnchoredPath -> Blob m -> Blob m -> a) -> Tree m -> Tree m -> [a]
zipCommonFiles f a b = catMaybes [ flip (f p) x `fmap` findFile a p
                                   | (p, File x) <- list b ]

-- | For each file in each of the two supplied trees, evaluate the supplied
-- function (supplying the corresponding file from the other tree, or Nothing)
-- and accumulate the results in a list. Hint: to get IO actions through, just
-- use sequence on the resulting list.  NB. This won't expand any stubs.
zipFiles :: (AnchoredPath -> Maybe (Blob m) -> Maybe (Blob m) -> a)
         -> Tree m -> Tree m -> [a]
zipFiles f a b = [ f p (findFile a p) (findFile b p)
                   | p <- paths a `sortedUnion` paths b ]
    where paths t = sort [ p | (p, File _) <- list t ]

zipTrees :: (AnchoredPath -> Maybe (TreeItem m) -> Maybe (TreeItem m) -> a)
         -> Tree m -> Tree m -> [a]
zipTrees f a b = [ f p (find a p) (find b p)
                   | p <- reverse (paths a `sortedUnion` paths b) ]
    where paths t = sort [ p | (p, _) <- list t ]

-- | Helper function for taking the union of AnchoredPath lists that
-- are already sorted.  This function does not check the precondition
-- so use it carefully.
sortedUnion :: [AnchoredPath] -> [AnchoredPath] -> [AnchoredPath]
sortedUnion [] ys = ys
sortedUnion xs [] = xs
sortedUnion a@(x:xs) b@(y:ys) = case compare x y of
                                LT -> x : sortedUnion xs b
                                EQ -> x : sortedUnion xs ys
                                GT -> y : sortedUnion a ys

-- | Cautiously extracts differing subtrees from a pair of Trees. It will never
-- do any unneccessary expanding. Tree hashes are used to cut the comparison as
-- high up the Tree branches as possible. The result is a pair of trees that do
-- not share any identical subtrees. They are derived from the first and second
-- parameters respectively and they are always fully expanded. It might be
-- advantageous to feed the result into 'zipFiles' or 'zipTrees'.
diffTrees :: forall m. (Functor m, Monad m) => Tree m -> Tree m -> m (Tree m, Tree m)
diffTrees left right =
            if treeHash left `match` treeHash right
               then return (emptyTree, emptyTree)
               else diff left right
  where isFile (File _) = True
        isFile _ = False
        notFile = not . isFile
        isEmpty = null . listImmediate
        subtree :: TreeItem m -> m (Tree m)
        subtree (Stub x _) = x
        subtree (SubTree x) = return x
        subtree (File _) = error "diffTrees tried to descend a File as a subtree"
        maybeUnfold (Stub x _) = SubTree `fmap` (x >>= expand)
        maybeUnfold (SubTree x) = SubTree `fmap` expand x
        maybeUnfold i = return i
        immediateN t = [ n | (n, _) <- listImmediate t ]
        diff left' right' = do
          is <- sequence [
                   case (lookup left' n, lookup right' n) of
                     (Just l, Nothing) -> do
                       l' <- maybeUnfold l
                       return (n, Just l', Nothing)
                     (Nothing, Just r) -> do
                       r' <- maybeUnfold r
                       return (n, Nothing, Just r')
                     (Just l, Just r)
                         | itemHash l `match` itemHash r ->
                             return (n, Nothing, Nothing)
                         | notFile l && notFile r ->
                             do x <- subtree l
                                y <- subtree r
                                (x', y') <- diffTrees x y
                                if isEmpty x' && isEmpty y'
                                   then return (n, Nothing, Nothing)
                                   else return (n, Just $ SubTree x', Just $ SubTree y')
                         | isFile l && isFile r ->
                             return (n, Just l, Just r)
                         | otherwise ->
                             do l' <- maybeUnfold l
                                r' <- maybeUnfold r
                                return (n, Just l', Just r')
                     _ -> error "n lookups failed"
                   | n <- immediateN left' `union` immediateN right' ]
          let is_l = [ (n, l) | (n, Just l, _) <- is ]
              is_r = [ (n, r) | (n, _, Just r) <- is ]
          return (makeTree is_l, makeTree is_r)

-- | Modify a Tree (by replacing, or removing or adding items).
modifyTree :: (Monad m) => Tree m -> AnchoredPath -> Maybe (TreeItem m) -> Tree m

modifyTree _ (AnchoredPath []) (Just (SubTree sub)) = sub

modifyTree t (AnchoredPath [n]) (Just item) =
    t { items = countmap items' `seq` items'
      , treeHash = NoHash }
  where !items' = M.insert n item (items t)

modifyTree t (AnchoredPath [n]) Nothing =
    t { items = countmap items' `seq` items'
      , treeHash = NoHash }
  where !items' = M.delete n (items t)

modifyTree t path@(AnchoredPath (n:r)) item =
    t { items = countmap items' `seq` items'
      , treeHash = NoHash }
  where subtree s = modifyTree s (AnchoredPath r) item
        !items' = M.insert n sub (items t)
        !sub = case lookup t n of
                Just (SubTree s) -> SubTree $! subtree s
                Just (Stub s _) -> Stub (do x <- s
                                            return $! subtree x) NoHash
                Nothing -> SubTree $! subtree emptyTree
                _ -> error $ "Modify tree at " ++ show path

modifyTree _ (AnchoredPath []) (Just (Stub _ _)) =
    error "Bug in descent in modifyTree."
modifyTree _ (AnchoredPath []) (Just (File _)) =
    error "Bug in descent in modifyTree."
modifyTree _ (AnchoredPath []) Nothing =
    error "Bug in descent in modifyTree."

countmap = M.fold (\_ i -> i + 1) 0

updateSubtrees :: (Tree m -> Tree m) -> Tree m -> Tree m
updateSubtrees fun t =
    fun $ t { items = M.mapWithKey (curry $ snd . update) $ items t
            , treeHash = NoHash }
  where update (k, SubTree s) = (k, SubTree $ updateSubtrees fun s)
        update (k, File f) = (k, File f)
        update (_, Stub _ _) = error "Stubs not supported in updateTreePostorder"

-- | Does /not/ expand the tree.
updateTree :: (Functor m, Monad m) => (TreeItem m -> m (TreeItem m)) -> Tree m -> m (Tree m)
updateTree fun t = do
    immediate <- mapM update $ listImmediate t
    SubTree t' <- fun . SubTree $ t { items = M.fromList immediate
                                    , treeHash = NoHash }
    return t'
  where update (k, SubTree tree) = (\new -> (k, SubTree new)) <$> updateTree fun tree
        update (k, item) = (\new -> (k, new)) <$> fun item

-- | Lay one tree over another. The resulting Tree will look like the base (1st
-- parameter) Tree, although any items also present in the overlay Tree will be
-- taken from the overlay. It is not allowed to overlay a different kind of an
-- object, nor it is allowed for the overlay to add new objects to base.  This
-- means that the overlay Tree should be a subset of the base Tree (although
-- any extraneous items will be ignored by the implementation).
overlay :: (Functor m, Monad m) => Tree m -> Tree m -> Tree m
overlay base over = Tree { items = M.fromList immediate
                         , treeHash = NoHash }
    where immediate = [ (n, get n) | (n, _) <- listImmediate base ]
          get n = case (M.lookup n $ items base, M.lookup n $ items over) of
                    (Just (File _), Just f@(File _)) -> f
                    (Just (SubTree b), Just (SubTree o)) -> SubTree $ overlay b o
                    (Just (Stub b _), Just (SubTree o)) -> Stub (flip overlay o `fmap` b) NoHash
                    (Just (SubTree b), Just (Stub o _)) -> Stub (overlay b `fmap` o) NoHash
                    (Just (Stub b _), Just (Stub o _)) -> Stub (do o' <- o
                                                                   b' <- b
                                                                   return $ overlay b' o') NoHash
                    (Just x, _) -> x
                    (_, _) -> error $ "Unexpected case in overlay at get " ++ show n ++ "."