arbtt-0.10: 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
import Data.Aeson (parseJSON, json)
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)))