arbtt-0.12.0.3: src/import-main.hs

{-# LANGUAGE CPP #-}
module Main where

import System.Directory
import System.FilePath
import System.Console.GetOpt
import System.Environment
import System.Exit
import System.IO
import qualified Data.Map as M
import Data.Version (showVersion)
import Data.Maybe
import Control.Monad
import Control.Applicative
import Data.Conduit.Attoparsec
import Data.Conduit
#if MIN_VERSION_aeson(2,2,0)
import Data.Aeson (parseJSON)
import Data.Aeson.Parser (json)
#else
import Data.Aeson (parseJSON, json)
#endif
import Data.Aeson.Types (parseEither)
import Data.Binary.StringRef
import Data.Attoparsec.ByteString.Char8 (skipSpace, option)
import qualified Data.ByteString.Char8 as BS
import qualified Data.ByteString.Lazy as BSL
import qualified Data.Conduit.List as C
import qualified Data.Conduit.Binary as C

import TimeLog
import Data
import CommonStartup
import DumpFormat
import LockFile

import Paths_arbtt (version)

data Options = Options
    { optLogFile :: String
    , optFormat :: DumpFormat
    , optAppend :: Bool
    }

defaultOptions dir = Options
    { optLogFile = dir </> "capture.log"
    , optAppend = False
    , optFormat = DFShow
    }

versionStr = "arbtt-import " ++ showVersion version
header = "Usage: arbtt-import [OPTIONS...]"

options :: [OptDescr (Options -> IO Options)]
options =
     [ Option "h?"     ["help"]
              (NoArg $ \_ -> do
                    putStr (usageInfo header options)
                    exitSuccess
              )
              "show this help"
     , Option "V"      ["version"]
              (NoArg $ \_ -> do
                    putStrLn versionStr
                    exitSuccess
              )
              "show the version number"
     , Option "f"      ["logfile"]
              (ReqArg (\arg opt -> return opt { optLogFile = arg }) "FILE")
               "use this file instead of ~/.arbtt/capture.log"
     , Option "a"      ["append"]
              (NoArg (\opt -> return opt { optAppend = True } ))
               "append to the logfile, instead of overriding it"
     , Option "t"      ["format"]
              (ReqArg (\arg opt ->
                case readDumpFormat arg of
                    Just DFHuman -> do
                        hPutStrLn stderr "The “Human” format cannot imported."
                        hPutStr stderr (usageInfo header options)
                        exitFailure
                    Just fm -> return $ opt { optFormat = fm}
                    Nothing -> do
                        hPutStrLn stderr ("Invalid format \"" ++ arg ++ "\".")
                        hPutStr stderr (usageInfo header options)
                        exitFailure
                    ) "FORMAT")
               "output format, one of Show (default) or JSON "
     ]

parseConduit :: DumpFormat -> Conduit BS.ByteString IO [TimeLogEntry CaptureData]
parseConduit DFHuman = error "Cannot read back human format"
parseConduit DFShow = C.lines =$= C.map ( (:[]) . read . BS.unpack)
parseConduit DFJSON =
        conduitParser (ignoreWhiteSpace json)
    =$= C.map snd
    =$= C.catMaybes
    =$= tlConduit
  where
    tlConduit = C.mapM $ \ v ->
        case parseEither oneOrMany v of
            Left e -> do
                hPutStrLn stderr ("Cannot parse log entry: " ++ e)
                exitFailure
            Right x -> pure x

    oneOrMany v = parseJSON v
            <|> ((:[]) <$> parseJSON v)

    ignoreWhiteSpace p = skipSpace *> option Nothing (Just <$> p)

binaryConduit :: ListOfStringable a =>
    Conduit [(TimeLogEntry a, Maybe a)] IO (Flush BS.ByteString)
binaryConduit = C.map (map go) =$= C.concatMap chunk
  where
    go (x,prev) = BSL.toStrict $ ls_encode strs x
        where strs = maybe [] listOfStrings prev
    chunk xs = map Chunk xs ++ [Flush]

main = do
  commonStartup
  args <- getArgs
  actions <- case getOpt Permute options args of
          (o,[],[])  -> return o
          (_,_,errs) -> do
                hPutStr stderr (concat errs ++ usageInfo header options)
                exitFailure

  dir <- getAppUserDataDirectory "arbtt"
  flags <- foldl (>>=) (return (defaultOptions dir)) actions

  lockFile (optLogFile flags)

  ex <- doesFileExist (optLogFile flags)
  when (ex && not (optAppend flags)) $ do
      putStrLn $ "File at " ++ (optLogFile flags) ++ " does already exist. Please delete this"
      putStrLn $ "file before running arbtt-import."
      exitFailure
  createTimeLog False (optLogFile flags)

  h <- openBinaryFile (optLogFile flags) AppendMode
  runConduit $ C.sourceHandle stdin
    =$= parseConduit (optFormat flags)
    =$= stutterList
    =$= binaryConduit
    =$= C.sinkHandleFlush h
  hClose h


stutter :: Monad m => Conduit (TimeLogEntry a) m (TimeLogEntry a, Maybe a)
stutter =
    loop Nothing
  where
    loop prev = do
        mx <- await
        case mx of
            Nothing -> return ()
            Just x -> do
                yield (x, prev)
                loop (Just (tlData x))

stutterList :: Monad m => Conduit [TimeLogEntry a] m [(TimeLogEntry a, Maybe a)]
stutterList =
    loop Nothing
  where
    loop prev = do
        mx <- await
        case mx of
            Nothing -> return ()
            Just xs -> do
                let xs' = zip xs (prev : map (Just . tlData) xs)
                yield xs'
                loop (Just (tlData (last xs)))