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

-- | A module for shell-like / perl-like programming in Haskell. The stuff in
-- here is not pretty, but it does get job done. The functionality provided by
-- this module is (unlike standard Haskell filesystem functionality)
-- thread-safe: each ShIO maintains its own environment and its own working
-- directory.
module Shellish
       (
         -- * Entering ShIO.
         ShIO, shellish, sub, silently, verbosely

         -- * Modifying and querying environment.
         , setenv, getenv, cd, pwd

         -- * Printing & stuff.
         , echo, echo_n, echo_err, echo_n_err

         -- * Querying filesystem.
         , ls, test_e, test_f, test_d, test_s, which, find

         -- * Manipulating filesystem.
         , mv, rm_f, rm_rf, cp, cp_r, mkdir, mkdir_p
         , readfile, writefile, appendfile, withTmpDir

         -- * Running external commands.
         , run, (#), lastOutput, lastStdout, lastStderr

         -- * Utilities.
         , (</>), (<.>), (<$>), (<$$>), grep, whenM, canonic
         , catch_sh, liftIO, MemTime(..), time, catchany
         , RunFailed(..)
         ) where

import Prelude hiding ( catch, readFile )
import Data.List( isInfixOf, (\\) )
import Data.Typeable
import Data.IORef
import Data.Maybe
import System.IO hiding ( readFile )
import System.IO.Strict( readFile )
import System.PosixCompat.Files( getSymbolicLinkStatus, isSymbolicLink )
import System.Directory
import System.Exit
import System.FilePath
import System.Environment
import Control.Applicative
import Control.Exception
import Control.Monad.Reader
import Control.Concurrent
import Data.Time.Clock( getCurrentTime, diffUTCTime, UTCTime(..) )
import Control.Exception ( bracket, evaluate )
import qualified Data.ByteString.Char8 as B
import System.Process( runInteractiveProcess, waitForProcess, ProcessHandle )

data MemTime = MemTime Rational Double deriving (Read, Show, Ord, Eq)

data St = St { sCode :: Int, sStderr :: B.ByteString , sStdout :: B.ByteString
             , sOutput :: B.ByteString, sDirectory :: FilePath
             , sVerbose :: Bool, sRun :: String -> [String] ->
                                         ShIO (Handle, Handle, Handle, ProcessHandle)
             , sEnvironment :: [(String, String)] }

type ShIO a = ReaderT (IORef St) IO a

get :: ShIO St
get = ask >>= liftIO . readIORef
put v = ask >>= liftIO . flip writeIORef v
modify f = ask >>= liftIO . flip modifyIORef f
gets f = f <$> get

runInteractiveProcess' cmd args = do
  st <- get
  liftIO $ runInteractiveProcess cmd args (Just $ sDirectory st) (Just $ sEnvironment st)

-- | A helper to catch any exception (same as
-- @... `catch` \(e :: SomeException) -> ...@).
catchany :: IO a -> (SomeException -> IO a) -> IO a
catchany = catch

-- | Catch an exception in the ShIO monad.
catch_sh :: (Exception e) => ShIO a -> (e -> ShIO a) -> ShIO a
catch_sh a h = do ref <- ask
                  liftIO $ catch (runReaderT a ref) (\e -> runReaderT (h e) ref)

-- | Change current working directory of ShIO. This does *not* change the
-- working directory of the process we are running it. Instead, ShIO keeps
-- track of its own workking directory and builds absolute paths internally
-- instead of passing down relative paths. This may have performance
-- repercussions if you are doing hundreds of thousands of filesystem
-- operations. You will want to handle these issues differently in those cases.
cd :: FilePath -> ShIO ()
cd dir = do dir' <- path dir
            modify $ \st -> st { sDirectory = dir' }

path p | isRelative p = (</> p) <$> gets sDirectory
       | otherwise = return p

-- | Currently a "renameFile" wrapper. TODO: Support cross-filesystem
-- move. TODO: Support directory paths in the second parameter, like in "cp".
mv :: FilePath -> FilePath -> ShIO ()
mv a b = do a' <- path a
            b' <- path b
            liftIO $ renameFile a' b'

-- | List directory contents. Does *not* include \".\" and \"..\", but it does
-- include (other) hidden files.
ls :: FilePath -> ShIO [String]
ls dir = do dir' <- path dir
            liftIO $ filter (`notElem` [".", ".."]) <$> getDirectoryContents dir'

-- | List directory recursively (like the POSIX utility "find").
find :: FilePath -> ShIO [String]
find dir = do bits <- ls dir
              sub <- forM bits $ \x -> do
                ex <- test_d $ dir </> x
                sym <- test_s $ dir </> x
                if ex && not sym then find (dir </> x)
                                 else return []
              return $ map (dir </>) bits ++ concat sub

-- | Obtain the current (ShIO) working directory.
pwd :: ShIO String
pwd = gets sDirectory

-- | Echo string to standard (error, when using _err variants) output. The _n
-- variants do not print a final newline.
echo, echo_n, echo_err, echo_n_err :: String -> ShIO ()
echo = liftIO . putStrLn
echo_n = liftIO . (>> hFlush System.IO.stdout) . putStr
echo_err = liftIO . hPutStrLn stderr
echo_n_err = liftIO . (>> hFlush stderr) . hPutStr stderr

-- | Create a new directory (fails if the directory exists).
mkdir :: FilePath -> ShIO ()
mkdir = path >=> liftIO . createDirectory

-- | Create a new directory, including parents (succeeds if the directory
-- already exists).
mkdir_p :: FilePath -> ShIO ()
mkdir_p = path >=> liftIO . createDirectoryIfMissing True

-- | Get a full path to an executable on @PATH@, if exists. FIXME does not
-- respect setenv'd environment and uses @PATH@ inherited from the process
-- environment.
which :: String -> ShIO (Maybe FilePath)
which = liftIO . findExecutable

-- | Obtain a (reasonably) canonic file path to a filesystem object. Based on
-- "canonicalizePath" in System.FilePath.
canonic :: FilePath -> ShIO FilePath
canonic = path >=> liftIO . canonicalizePath

-- | A monadic-conditional version of the "when" guard.
whenM :: Monad m => m Bool -> m () -> m ()
whenM c a = do res <- c
               when res a

-- | Does a path point to an existing filesystem object?
test_e :: FilePath -> ShIO Bool
test_e f = do f' <- path f
              liftIO $ do
                dir <- doesDirectoryExist f'
                file <- doesFileExist f'
                return $ file || dir

-- | Does a path point to an existing file?
test_f :: FilePath -> ShIO Bool
test_f = path >=> liftIO . doesFileExist

-- | Does a path point to an existing directory?
test_d :: FilePath -> ShIO Bool
test_d = path >=> liftIO . doesDirectoryExist

-- | Does a path point to a symlink?
test_s :: FilePath -> ShIO Bool
test_s = path >=> liftIO . \f -> do
  stat <- getSymbolicLinkStatus f
  return $ isSymbolicLink stat

-- | A swiss army cannon for removing things. Actually this goes farther than a
-- normal rm -rf, as it will circumvent permission problems for the files we
-- own. Use carefully.
rm_rf :: FilePath -> ShIO ()
rm_rf f = path f >>= \f' -> do
  let deletable = Permissions True True True True
  whenM (test_d f) $ do
    find f' >>= mapM (\file -> liftIO $ setPermissions file deletable `catchany` \_ -> return ())
    liftIO $ removeDirectoryRecursive f'
  whenM (test_f f) $ rm_f f'

-- | Remove a file. Does not fail if the file already is not there. Does fail
-- if the file is not a file.
rm_f :: FilePath -> ShIO ()
rm_f f = path f >>= \f' -> whenM (test_e f) $ liftIO $ removeFile f'

-- | Set an environment variable. The environment is maintained in ShIO
-- internally, and is passed to any external commands to be executed.
setenv :: String -> String -> ShIO ()
setenv k v = modify $ \x -> x { sEnvironment = wibble $ sEnvironment x }
  where wibble env = (k, v) : filter ((/=k).fst) env

-- | Fetch the current value of an environment variable. Both empty and
-- non-existent variables give empty string as a result.
getenv :: String -> ShIO String
getenv k = fromMaybe "" <$> lookup k <$> gets sEnvironment

-- | Create a sub-ShIO in which external command outputs are not echoed. See "sub".
silently :: ShIO a -> ShIO a
silently a = sub $ modify (\x -> x { sVerbose = False }) >> a

-- | Create a sub-ShIO in which external command outputs are echoed. See "sub".
verbosely :: ShIO a -> ShIO a
verbosely a = sub $ modify (\x -> x { sVerbose = True }) >> a

-- | Enter a sub-ShIO. The new ShIO inherits the environment and working
-- directory from the current one, but the sub-ShIO cannot affect the current
-- one. Exceptions are propagated normally.
sub :: ShIO a -> ShIO a
sub a = do
  st <- get
  r <- a `catch_sh` (\(e :: SomeException) -> put st >> throw e)
  put st
  return r

-- | Enter a ShIO from (Monad)IO. The environment and working directories are
-- inherited from the current process-wide values. Any subsequent changes in
-- processwide working directory or environment are not reflected in the
-- running ShIO.
shellish :: MonadIO m => ShIO a -> m a
shellish a = do
  env <- liftIO $ getEnvironment
  dir <- liftIO $ getCurrentDirectory
  let empty = St { sCode = 0, sStderr = B.empty, sOutput = B.empty
                 , sStdout = B.empty, sVerbose = True
                 , sRun = runInteractiveProcess', sEnvironment = env
                 , sDirectory = dir }
  stref <- liftIO $ newIORef empty
  liftIO $ runReaderT a stref

drain :: Handle -> Maybe Handle -> Chan B.ByteString -> ShIO (Chan B.ByteString)
drain h verb all =
  do chan <- liftIO newChan
     let work acc = do line <- liftIO $ B.hGetLine h
                       writeChan all line
                       liftIO $ when (isJust verb) $ B.hPutStrLn (fromJust verb) line
                       work $ B.concat [acc, line, "\n"]
                     `catchany` \_ -> liftIO $ writeChan chan acc
     _ <- liftIO $ forkIO $ work ""
     return chan

drainChan :: Chan a -> IO [a]
drainChan ch = do empty <- isEmptyChan ch
                  if empty then return []
                           else do b <- readChan ch
                                   (b:) <$> drainChan ch

data RunFailed = RunFailed String Int String deriving (Typeable)

instance Show RunFailed where
  show (RunFailed cmd code errs) =
    "error running " ++ cmd ++ ": exit status " ++ show code ++ ":\n" ++ errs

instance Exception RunFailed

-- | An infix shorthand for "run". Write @\"command\" # [ \"argument\" ... ]@.
(#) :: String -> [String] -> ShIO String
cmd # args = run cmd args

-- | Execute an external command. Takes the command name (no shell allowed,
-- just a name of something that can be found via @PATH@; FIXME: setenv'd
-- @PATH@ is not taken into account, only the one inherited from the actual
-- outside environment). Nothing is provided on "stdin" of the process, and
-- "stdout" and "stderr" are collected and stored. The "stdout" is returned as
-- a result of "run", and complete outputs are available after the fact using
-- "lastStdout", "lastStderr" and "lastOutput" with the last giving an
-- interleaving of both, approximately reflecting the times of their arrival --
-- basically what @2>&1@ would give you in a shell.
run :: String -> [String] -> ShIO String
run cmd args = do
    st <- get
    (_,outH,errH,procH) <- (sRun st) cmd args
    all' <- liftIO $ newChan
    res' <- drain outH (if sVerbose st then Just stdout else Nothing) all'
    errs' <- drain errH (if sVerbose st then Just stderr else Nothing) all'
    ex <- liftIO $ waitForProcess procH
    errs <- liftIO $ readChan errs'
    res <- liftIO $ readChan res'
    all <- liftIO $ B.intercalate "\n" <$> drainChan all'
    modify $ \x -> x { sCode = 0, sStderr = errs, sStdout = res, sOutput = all }
    case ex of
      ExitSuccess -> return ()
      ExitFailure n -> do
        modify $ \x -> x { sCode = n }
        throw $ RunFailed (cmd ++ " " ++ show args) n (B.unpack errs)
    return $ B.unpack res

-- | The output of last external command. See "run".
lastOutput, lastStderr, lastStdout :: ShIO B.ByteString
lastStdout = gets sStdout
lastStderr = gets sStderr
lastOutput = gets sOutput

-- | Run a ShIO computation and collect timing (TODO: and memory) information.
time :: ShIO a -> ShIO (MemTime, a)
time what = sub $ do -- TODO track memory usage as well
  t <- liftIO getCurrentTime
  res <- what
  t' <- liftIO getCurrentTime
  let mt = MemTime 0 (realToFrac $ diffUTCTime t' t)
  return (mt, res)

{-
    stats_f <- liftIO $
      do tmpdir <- getTemporaryDirectory
         (f, h) <- openTempFile tmpdir "darcs-stats-XXXX"
         hClose h
         return f
    let args = args' ++ ["+RTS", "-s" ++ stats_f, "-RTS"]
    ...
    stats <- liftIO $ do c <- readFile' stats_f
                         removeFile stats_f `catchany` \e -> hPutStrLn stderr (show e)
                         return c
                       `catchany` \_ -> return ""
    let bytes = (stats =~ "([0-9, ]+) M[bB] total memory in use") :: String
        mem = case length bytes of
          0 -> 0
          _ -> (read (filter (`elem` "0123456789") bytes) :: Int)
    recordMemoryUsed $ mem * 1024 * 1024
    return res
-}

-- | Copy a file, or a directory recursively.
cp_r :: FilePath -> FilePath -> ShIO ()
cp_r from to = do
    whenM (test_d from) $
      mkdir to >> ls from >>= mapM_ (\item -> cp_r (from </> item) (to </> item))
    whenM (test_f from) $ cp from to

-- | Copy a file. The second path could be a directory, in which case the
-- original file name is used, in that directory.
cp :: FilePath -> FilePath -> ShIO ()
cp from to = do
  from' <- path from
  to' <- path to
  to_dir <- test_d to
  liftIO $ copyFile from' (if to_dir then to' </> takeFileName from else to')

class PredicateLike pattern hay where
  match :: pattern -> hay -> Bool

instance PredicateLike (a -> Bool) a where
  match = id

instance (Eq a) => PredicateLike [a] [a] where
  match pat = (pat `isInfixOf`)

-- | Like filter, but more conveniently used with String lists, where a
-- substring match (TODO: also provide regexps, and maybe globs) is expressed as
--  @grep \"needle\" [ \"the\", \"stack\", \"of\", \"hay\" ]@. Boolean
-- predicates just like with "filter" are supported too:
-- @grep (\"fun\" `isPrefixOf`) [...]@.
grep :: (PredicateLike pattern hay) => pattern -> [hay] -> [hay]
grep p l = filter (match p) l

-- | A functor-lifting function composition.
(<$$>) :: (Functor m) => (b -> c) -> (a -> m b) -> a -> m c
f <$$> v = fmap f . v

-- | Create a temporary directory and pass it as a parameter to a ShIO
-- computation. The directory is nuked afterwards.
withTmpDir :: (FilePath -> ShIO a) -> ShIO a
withTmpDir act = do
  dir <- liftIO $ getTemporaryDirectory
  (path, _) <- liftIO $ openTempFile dir "tmp"
  rm_f path
  mkdir path
  a <- act path `catch_sh` \(e :: SomeException) -> rm_rf path >> throw e
  rm_rf path
  return a

-- | Write a String to a file.
writefile :: FilePath -> String -> ShIO ()
writefile f bits = path f >>= \f' -> liftIO (writeFile f' bits)

-- | Append a String to a file.
appendfile :: FilePath -> String -> ShIO ()
appendfile f bits = path f >>= \f' -> liftIO (appendFile f' bits)

-- | (Strictly) read file into a String.
readfile :: FilePath -> ShIO String
readfile = path >=> liftIO . readFile