{-# LANGUAGE ExistentialQuantification, ViewPatterns #-}

module Main (main) where

import Control.Arrow
import Control.Applicative
import Control.Monad

import System.Environment
import System.FilePath
import System.IO
import System.Directory

data Unfold n a = forall s. U s (s -> IO (Either a (n, [s])))

instance Functor (Unfold n) where
  fmap f (U seed step) = U seed step'
   where
   step' seed' = (f +++ id) <$> step seed'

instance Monad (Unfold n) where
  return a = U () (\_ -> return (Left a))
  U seed step >>= f = U (Left seed) st
   where
   st (Left sd) = step sd >>= \e -> case e of
     Left a -> st . Right $ f a
     Right (n, ss) -> return $ Right (n, map Left ss)
   st (Right (U seed' step')) = step' seed' >>= \e -> case e of
     Left a        -> return $ Left a
     Right (n, ss) -> return . Right $ (n, map (Right . flip U step') ss)

branch :: Unfold n a -> (a -> (n, [Unfold n b])) -> Unfold n b
branch (U seed step) f = U (Left seed) st
 where
 st (Left sd) = step sd >>= \e -> case e of
   Left a -> let (n, u) = f a in return $ Right (n, map Right u)
   Right (n, ss) -> return $ Right (n, map Left ss)
 st (Right (U seed' step')) = step' seed' >>= \e -> case e of
   Left b        -> return $ Left b
   Right (n, ss) -> return . Right $ (n, map (Right . flip U step') ss)

liftU :: IO a -> Unfold n a
liftU ioa = U () (\_ -> Left <$> ioa)


data Fold n a = forall r. F (a -> IO r) (n -> [r] -> IO r)

find :: FilePath -> Unfold FilePath FilePath
find = flip U g
 where
 p "."  = False
 p ".." = False
 p _    = True
 g f = do b <- doesDirectoryExist f
          if b
             then Right . (,) f . map (f </>) . filter p <$> getDirectoryContents f
             else return $ Left f

find' :: FilePath -> Unfold FilePath FilePath
find' f = do b <- liftU $ doesDirectoryExist f
             if b
               then liftU (getDirectoryContents f) `branch` \fs ->
                    (f, map (find' . (f </>)) . filter p $ fs)
               else return f
 where
 p "."  = False
 p ".." = False
 p _    = True

sizes :: FilePath -> Unfold FilePath (FilePath, Integer)
sizes = find' >=> \f' -> liftU $ do h <- openFile f' ReadMode
                                    i <- hFileSize h
                                    hClose h
                                    return (f', i)
                                       
annihilate :: Unfold n a -> Fold n a -> IO ()
annihilate (U seed step) (F leaf node) = go seed >> return ()
 where
 go s = step s >>= \v -> case v of
                              Left a        -> leaf a
                              Right (n, ss) -> mapM go ss >>= node n

processDir f (sum -> n) = do putStr $ f ++ ": "
                             print n
                             return n

processFile (f, n) = do putStr $ f ++ ": "
                        print n
                        return n

process = F processFile processDir

main = do (f:_) <- getArgs
          annihilate (sizes f) process