-----------------------------------------------------------------------------
-- |
-- Module      :  System.Directory.DirTree
-- Copyright   :  Suite Solutions Ltd., Israel 2010
--
-- Maintainer  :  Yitzchak Gale <yitzg@suite-sol.com>
-- Portability :  portable
--
-- Creates a lazy tree in the shape of a filesystem directory
-- tree. The labels of the tree are the paths to the files and
-- directories. Functions and class instances from "Data.Tree" and
-- "Data.Tree.Utils" can then be used to build powerful
-- combinator-style operations on the filesystem.

{- Copyright (c) 2010 Suite Solutions Ltd., Israel. All rights reserved.
For licensing information, see the BSD3-style license in the file
LICENSE that was originally distributed by the author together with
this file. -}

module System.Directory.DirTree where

import Data.Tree
import System.Directory
import System.FilePath
import Data.Functor ((<$>))
import System.IO.Unsafe (unsafeInterleaveIO)

-- | An entry in a directory is a path to a file or a directory.
data FileOrDir path = File { fdPath :: path } | Dir { fdPath :: path }
  deriving (Eq, Ord, Show)

instance Functor FileOrDir where
  fmap f (File path) = File $ f path
  fmap f (Dir  path) = Dir  $ f path

-- | A tree that represents a directory structure
type DirTree = Tree (FileOrDir FilePath)

-- | Determine whether a directory entry represents a directory.
isDir :: FileOrDir a -> Bool
isDir (Dir _) = True
isDir _       = False

-- | Get the directory tree rooted at the given path, with each
-- entry a relative path to the root.
getRelDirTree :: FilePath -> IO DirTree
getRelDirTree root = unsafeLazyUnfoldTree build base
  where
    (dir, base) = splitFileName root
    build relPath = do
      let path = dir </> relPath
      dir <- doesDirectoryExist path
      if dir
        then (,) (Dir relPath) . map (relPath </>) <$> safeDirContents path
        else return (File relPath, [])
    safeDirContents = fmap (filter $ not . dots) . getDirectoryContents
    dots = all (== '.') . take 2

-- | Get the directory tree rooted at the given path, with each
-- entry a name of a file or directory.
getFileNameDirTree :: FilePath -> IO DirTree
getFileNameDirTree = (fmap . fmap . fmap) takeFileName . getRelDirTree

-- | Get the directory tree rooted at the given path, with each
-- entry an extension of the original root path
getDirTree :: FilePath -> IO DirTree
getDirTree root = (fmap . fmap . fmap) (dir </>) $ getRelDirTree root
  where dir = takeDirectory root

-- | Build a tree using lazy IO. This is a lazy version
-- of 'Data.Tree.unfoldTreeM'.
unsafeLazyUnfoldTree :: (b -> IO (a, [b])) -> b -> IO (Tree a)
unsafeLazyUnfoldTree build seed = do
  (label, seeds) <- build seed
  Node label <$> mapM (unsafeInterleaveIO . unsafeLazyUnfoldTree build) seeds