content: {-# LANGUAGE CPP                       #-}
  {-# LANGUAGE FlexibleContexts          #-}
  {-# LANGUAGE MultiParamTypeClasses     #-}
  {-# LANGUAGE NoMonomorphismRestriction #-}
  {-# LANGUAGE OverloadedStrings         #-}
  {-# LANGUAGE RankNTypes                #-}
  {-# LANGUAGE RecordWildCards           #-}
  {-# LANGUAGE TemplateHaskell           #-}
  
  -----------------------------------------------------------------------------
  -- |
  -- Module      :
  -- Copyright   :
  -- License     :
  --
  -- Maintainer  :
  -- Stability   :  experimental
  --
  -- Emulate all hadoop operations locally
  ----------------------------------------------------------------------------
  
  module Hadron.Run.Local where
  
  -------------------------------------------------------------------------------
  import           Control.Applicative          as A
  import           Control.Error
  import           Control.Lens
  import           Control.Monad.Reader
  import           Control.Monad.Trans.Resource
  import qualified Data.ByteString.Char8        as B
  import           Data.Conduit
  import           Data.Conduit.Binary          (sourceFile)
  import           Data.Default
  import           Data.Hashable
  import           Data.List
  import           Data.Monoid
  import           Data.Text                    (Text)
  import qualified Data.Text                    as T
  import           System.Directory
  import           System.Environment
  import           System.Exit
  import           System.FilePath.Glob
  import           System.FilePath.Lens
  import           System.FilePath.Posix
  import           System.IO
  import           System.Process
  -------------------------------------------------------------------------------
  import           Hadron.Logger
  import           Hadron.Run.FanOut
  import qualified Hadron.Run.Hadoop            as H
  import           Hadron.Utils
  -------------------------------------------------------------------------------
  #if MIN_VERSION_base(4, 7, 0)
  #else
  import           System.Posix.Env
  #endif
  
  
  
  newtype LocalFile = LocalFile { _unLocalFile :: FilePath }
      deriving (Eq,Show,Read,Ord)
  makeLenses ''LocalFile
  
  
  data LocalRunSettings = LocalRunSettings {
        _lrsTempPath :: FilePath
      -- ^ Root of the "file system" during a localrun
      }
  makeLenses ''LocalRunSettings
  
  
  instance Default LocalRunSettings where
      def = LocalRunSettings "tmp"
  
  
  type Local = ReaderT LocalRunSettings IO
  
  runLocal :: r -> ReaderT r m a -> m a
  runLocal env f = runReaderT f env
  
  
  -------------------------------------------------------------------------------
  path :: (MonadIO m, MonadReader LocalRunSettings m) => LocalFile -> m FilePath
  path (LocalFile fp) = do
      root <- view lrsTempPath
      let p = root </> fp
          dir = p ^. directory
      liftIO $ createDirectoryIfMissing True dir
      return p
  
  
  -------------------------------------------------------------------------------
  getRecursiveDirectoryContents :: FilePath -> IO [FilePath]
  getRecursiveDirectoryContents dir0 = go dir0
      where
        go dir = do
            fs <- liftM (map (dir </>) . filter (not . flip elem [".", ".."])) $
              getDirectoryContents' dir
            fs' <- filterM (fmap not . doesDirectoryExist) fs
            fss <- mapM go fs
            return $ fs' ++ concat fss
  
  
  -------------------------------------------------------------------------------
  -- | A version that return [] instead of an error when directory does not exit.
  getDirectoryContents' :: FilePath -> IO [FilePath]
  getDirectoryContents' fp = do
      chk <- doesDirectoryExist fp
      case chk of
        False -> return []
        True -> getDirectoryContents fp
  
  
  -------------------------------------------------------------------------------
  -- | Recursive contents if given a directory, assumed glob pattern if not.
  getInputFiles :: FilePath -> IO [FilePath]
  getInputFiles fp = do
      chk <- doesDirectoryExist fp
      case chk of
        True -> glob (fp </> "**")
        False -> glob fp
  
  
  -------------------------------------------------------------------------------
  localMapReduce
      :: (Functor m, MonadIO m)
      => LocalRunSettings
      -> String                   -- ^ MapReduceKey
      -> String                   -- ^ RunToken
      -> H.HadoopRunOpts
      -> ExceptT Text m ()
  localMapReduce lrs mrKey token H.HadoopRunOpts{..} = do
      exPath <- scriptIO getExecutablePath
      echoInfo $ "Launching Hadoop job for MR key: " <> ls mrKey
  
  
      expandedInput <- liftIO $ liftM concat $ forM _mrsInput $ \ inp ->
        withLocalFile lrs (LocalFile inp) getInputFiles
  
  
      let enableCompress = case _mrsCompress of
            Nothing -> False
            Just x -> isInfixOf "Gzip" x
  
          -- Are the input files already compressed?
          inputCompressed file = isInfixOf ".gz" file
  
  
      outFile <- liftIO $ withLocalFile lrs (LocalFile _mrsOutput) $ \ fp ->
        case fp ^. extension . to null of
          False -> return fp
          True -> do
            createDirectoryIfMissing True fp
            return $ fp </> ("0000.out" ++ if enableCompress then ".gz" else "")
  
  
      let pipe = " | "
  
          maybeCompress = if enableCompress
                          then  pipe <> "gzip"
                          else ""
  
          maybeGunzip fp = (if inputCompressed fp then ("gunzip" <> pipe) else "")
  
          maybeReducer = case _mrsNumReduce of
            Just 0 -> ""
            _ -> pipe <> exPath <> " " <> token <> " " <> "reducer_" <> mrKey
  
  
  
          -- map over each file individually and write results into a temp file
          mapFile infile = clearExit . scriptIO . withTmpMapFile infile $ \ fp -> do
              echoInfo ("Running command: " <> ls (command fp))
  #if MIN_VERSION_base(4, 7, 0)
              setEnv "mapreduce_map_input_file" infile
  #else
              setEnv "mapreduce_map_input_file" infile True
  #endif
              system (command fp)
            where
              command fp =
                  "cat " <> infile <> pipe <>
                  maybeGunzip infile <>
                  exPath <> " " <> token <> " " <> "mapper_" <> mrKey <>
                  " > " <> fp
  
  
          -- a unique temp file for each input file
          withTmpMapFile infile f = liftIO $
            withLocalFile lrs (LocalFile ((show (hash infile)) <> "_mapout")) f
  
  
          getTempMapFiles = mapM (flip withTmpMapFile return) expandedInput
  
          -- concat all processed map output, sort and run through the reducer
          reduceFiles = do
              fs <- getTempMapFiles
              echoInfo ("Running command: " <> ls (command fs))
              scriptIO $ createDirectoryIfMissing True (outFile ^. directory)
              clearExit $ scriptIO $ system (command fs)
            where
              command fs =
                  "cat " <> intercalate " " fs <> pipe <>
                  ("sort -t$'\t' -k1," <> show (H.numSegs _mrsPart)) <>
                  maybeReducer <>
                  maybeCompress <>
                  " > " <> outFile
  
          removeTempFiles = scriptIO $ do
              fs <- getTempMapFiles
              mapM_ removeFile fs
  
  
      echoInfo "Mapping over individual local input files."
      mapM_ mapFile expandedInput
  
      echoInfo "Executing reduce stage."
      reduceFiles
  
      removeTempFiles
  
  
  -------------------------------------------------------------------------------
  echo :: (A.Applicative m, MonadIO m) => Severity -> LogStr -> m ()
  echo sev msg = runLog $ logMsg "Local" sev msg
  
  
  -------------------------------------------------------------------------------
  echoInfo :: (Applicative m, MonadIO m) => LogStr -> m ()
  echoInfo = echo InfoS
  
  
  -------------------------------------------------------------------------------
  -- | Fail if command not successful.
  clearExit
      :: (Functor m, MonadIO m)
      => ExceptT Text m ExitCode
      -> ExceptT Text m ()
  clearExit f = do
      res <- f
      case res of
        ExitSuccess -> echoInfo "Command successful."
        e -> do
          echo ErrorS $ ls $ "Command failed: " ++ show e
          hoistEither $ Left $ T.pack $ "Command failed with: " ++ show e
  
  
  -------------------------------------------------------------------------------
  -- | Check if the target file is present.
  hdfsFileExists
      :: (MonadIO m, MonadReader LocalRunSettings m)
      => LocalFile
      -> m Bool
  hdfsFileExists p = liftIO . chk =<< path p
      where
        chk fp = (||) <$> doesFileExist fp <*> doesDirectoryExist fp
  
  
  -------------------------------------------------------------------------------
  hdfsDeletePath
      :: (MonadIO m, MonadReader LocalRunSettings m)
      => LocalFile
      -> m ()
  hdfsDeletePath p = do
      fp <- path p
      liftIO $ do
        chk <- doesDirectoryExist fp
        when chk (removeDirectoryRecursive fp)
        chk2 <- doesFileExist fp
        when chk2 (removeFile fp)
  
  
  -------------------------------------------------------------------------------
  hdfsLs
      :: (MonadIO m, MonadReader LocalRunSettings m)
      => LocalFile -> m [File]
  hdfsLs p = do
      fs <- liftIO . getDirectoryContents' =<< path p
      return $ map (File "" 1 "" "") $ map (_unLocalFile p </>) fs
  
  
  
  -------------------------------------------------------------------------------
  hdfsPut
      :: (MonadIO m, MonadReader LocalRunSettings m)
      => LocalFile
      -> LocalFile
      -> m ()
  hdfsPut src dest = do
      src' <- path src
      dest' <- path dest
      liftIO $ copyFile src' dest'
  
  
  -------------------------------------------------------------------------------
  -- | Create a new multiple output file manager.
  hdfsFanOut
      :: (MonadIO m, MonadReader LocalRunSettings m)
      => FilePath
      -- ^ Temporary file location.
      -> m FanOut
  hdfsFanOut tmp = do
      env <- ask
      liftIO $ mkFanOut (mkHandle env)
    where
  
        mkTmp fp = tmp </> fp
  
        -- write into a temp file loc until we know the stage is
        -- complete without failure
        mkHandle env fp = do
          fp' <- runLocal env $ path (LocalFile (mkTmp fp))
          createDirectoryIfMissing True (fp' ^. directory)
          h <- openFile fp' AppendMode
          return (h, fin env fp)
  
        -- move temp file to its final destination
        fin env fp0 = do
          temp <- runLocal env $ path (LocalFile (mkTmp fp0))
          dest <- runLocal env $ path (LocalFile fp0)
          createDirectoryIfMissing True (dest ^. directory)
          renameFile temp dest
  
  
  -------------------------------------------------------------------------------
  hdfsMkdir
      :: (MonadIO m, MonadReader LocalRunSettings m)
      => LocalFile
      -> m ()
  hdfsMkdir p = liftIO . createDirectoryIfMissing True =<< path p
  
  
  -------------------------------------------------------------------------------
  hdfsCat :: LocalFile -> Producer (ResourceT Local) B.ByteString
  hdfsCat p = sourceFile =<< (lift . lift) (path p)
  
  
  -------------------------------------------------------------------------------
  hdfsGet
      :: (MonadIO m, MonadReader LocalRunSettings m)
      => LocalFile
      -> m LocalFile
  hdfsGet fp = do
      target <- randomFileName
      hdfsPut fp target
      return target
  
  
  hdfsLocalStream :: LocalFile -> Producer (ResourceT Local) B.ByteString
  hdfsLocalStream = hdfsCat
  
  
  -------------------------------------------------------------------------------
  randomFileName :: MonadIO m => m LocalFile
  randomFileName = LocalFile `liftM` liftIO (randomToken 64)
  
  
  
  -------------------------------------------------------------------------------
  -- | Helper to work with relative paths using Haskell functions like
  -- 'readFile' and 'writeFile'.
  withLocalFile
      :: MonadIO m
      => LocalRunSettings
      -> LocalFile
      -- ^ A relative path in our working folder
      -> (FilePath -> m b)
      -- ^ What to do with the absolute path.
      -> m b
  withLocalFile settings fp f = f =<< runLocal settings (path fp)
  