Common.hs

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE QuasiQuotes #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Hledger.Flow.Common where

import Control.Concurrent.STM
import Control.Exception (IOException, try)
import qualified Control.Foldl as Fold
import Data.Char (isDigit)
import Data.Either
import Data.Function (on)
import qualified Data.List as List (groupBy, null, sortBy)
import Data.Maybe (mapMaybe)
import qualified Data.Map.Strict as Map
import Data.Ord (comparing)
import qualified Data.Text as T
import qualified Data.Text.IO as T
import qualified Data.Text.Read as T
import qualified GHC.IO.Handle.FD as H
import Hledger.Flow.BaseDir (relativeToBase, turtleBaseDir)
import Hledger.Flow.Logging
import Hledger.Flow.PathHelpers (AbsFile, TurtlePath, fromTurtleAbsFile, pathToTurtle)
import qualified Hledger.Flow.RuntimeOptions as RT
import Hledger.Flow.Types
import Path (absfile, relfile)
import qualified Path.IO as Path
import Turtle ((%), (<.>), (</>))
import qualified Turtle
import Prelude hiding (putStrLn, readFile, writeFile)

hledgerPathFromOption :: Maybe TurtlePath -> IO AbsFile
hledgerPathFromOption pathOption = do
  case pathOption of
    Just h -> do
      hlAbs <- fromTurtleAbsFile h
      isOnDisk <- Path.doesFileExist hlAbs
      if isOnDisk
        then return hlAbs
        else do
          let msg = Turtle.format ("Unable to find hledger at " % Turtle.fp) h
          errExit' 1 (T.hPutStrLn H.stderr) msg hlAbs
    Nothing -> do
      maybeH <- Path.findExecutable [relfile|hledger|]
      case maybeH of
        Just h -> return h
        Nothing -> do
          let msg =
                "Unable to find hledger in your path.\n"
                  <> "You need to either install hledger, or add it to your PATH, or provide the path to an hledger executable.\n\n"
                  <> "There are a number of installation options on the hledger website: https://hledger.org/download.html"
          errExit' 1 (T.hPutStrLn H.stderr) msg [absfile|/hledger|]

hledgerVersionFromPath :: TurtlePath -> IO T.Text
hledgerVersionFromPath hlp = fmap (T.strip . Turtle.linesToText) (Turtle.single $ shellToList $ Turtle.inproc (Turtle.format Turtle.fp hlp) ["--version"] Turtle.empty)

hledgerInfoFromPath :: Maybe TurtlePath -> IO HledgerInfo
hledgerInfoFromPath pathOption = do
  hlp <- hledgerPathFromOption pathOption
  hlv <- hledgerVersionFromPath $ pathToTurtle hlp
  return $ HledgerInfo hlp hlv

hledgerConfArgs :: RT.RuntimeOptions -> [T.Text]
hledgerConfArgs opts =
  case RT.hledgerConf opts of
    Nothing -> []
    Just conf -> ["--conf", Turtle.format Turtle.fp conf]

showCmdArgs :: [T.Text] -> T.Text
showCmdArgs args = T.intercalate " " (map escapeArg args)

escapeArg :: T.Text -> T.Text
escapeArg a = if T.count " " a > 0 then "'" <> a <> "'" else a

errExit :: Int -> TChan LogMessage -> T.Text -> a -> IO a
errExit exitStatus ch = errExit' exitStatus (channelErrLn ch)

errExit' :: Int -> (T.Text -> IO ()) -> T.Text -> a -> IO a
errExit' exitStatus logFun errorMessage dummyReturnValue = do
  logFun errorMessage
  Turtle.sleep 0.1
  _ <- Turtle.exit $ Turtle.ExitFailure exitStatus
  return dummyReturnValue

descriptiveOutput :: T.Text -> T.Text -> T.Text
descriptiveOutput outputLabel outTxt = do
  if not (T.null outTxt)
    then Turtle.format (Turtle.s % ":\n" % Turtle.s % "\n") outputLabel outTxt
    else ""

logTimedAction ::
  (HasVerbosity o) =>
  o ->
  TChan LogMessage ->
  T.Text ->
  [T.Text] ->
  T.Text ->
  [T.Text] ->
  (TChan LogMessage -> T.Text -> IO ()) ->
  (TChan LogMessage -> T.Text -> IO ()) ->
  IO FullOutput ->
  IO FullTimedOutput
logTimedAction opts ch cmdLabel extraCmdLabels cmd args stdoutLogger stderrLogger action = do
  logVerbose opts ch $ Turtle.format ("Begin: " % Turtle.s) cmdLabel
  logVerbose opts ch $ Turtle.format (Turtle.s % " " % Turtle.s) cmd (T.intercalate " " args)
  if (List.null extraCmdLabels) then return () else logVerbose opts ch $ T.intercalate "\n" extraCmdLabels
  timed@((ec, stdOut, stdErr), diff) <- Turtle.time action
  stdoutLogger ch stdOut
  stderrLogger ch stdErr
  logVerbose opts ch $ Turtle.format ("End:   " % Turtle.s % " " % Turtle.s % " (" % Turtle.s % ")") cmdLabel (Turtle.repr ec) (Turtle.repr diff)
  return timed

timeAndExitOnErr ::
  (HasSequential o, HasVerbosity o) =>
  o ->
  TChan LogMessage ->
  T.Text ->
  (TChan LogMessage -> T.Text -> IO ()) ->
  (TChan LogMessage -> T.Text -> IO ()) ->
  ProcFun ->
  ProcInput ->
  IO FullTimedOutput
timeAndExitOnErr opts ch cmdLabel = timeAndExitOnErr' opts ch cmdLabel []

timeAndExitOnErr' ::
  (HasSequential o, HasVerbosity o) =>
  o ->
  TChan LogMessage ->
  T.Text ->
  [T.Text] ->
  (TChan LogMessage -> T.Text -> IO ()) ->
  (TChan LogMessage -> T.Text -> IO ()) ->
  ProcFun ->
  ProcInput ->
  IO FullTimedOutput
timeAndExitOnErr' opts ch cmdLabel extraCmdLabels stdoutLogger stderrLogger procFun (cmd, args, stdInput) = do
  let action = procFun cmd args stdInput
  timed@((ec, stdOut, stdErr), _) <- logTimedAction opts ch cmdLabel extraCmdLabels cmd args stdoutLogger stderrLogger action
  case ec of
    Turtle.ExitFailure i -> do
      let cmdText = Turtle.format (Turtle.s % " " % Turtle.s) cmd $ showCmdArgs args
      let msgOut = descriptiveOutput "Standard output" stdOut
      let msgErr = descriptiveOutput "Error output" stdErr

      let exitMsg =
            Turtle.format
              ( "\n=== Begin Error: "
                  % Turtle.s
                  % " ===\nExternal command:\n"
                  % Turtle.s
                  % "\nExit code "
                  % Turtle.d
                  % "\n"
                  % Turtle.s
                  % Turtle.s
                  % "=== End Error: "
                  % Turtle.s
                  % " ===\n"
              )
              cmdLabel
              cmdText
              i
              msgOut
              msgErr
              cmdLabel
      errExit i ch exitMsg timed
    Turtle.ExitSuccess -> return timed

procWithEmptyOutput :: ProcFun
procWithEmptyOutput cmd args stdinput = do
  ec <- Turtle.proc cmd args stdinput
  return (ec, T.empty, T.empty)

parAwareProc :: (HasSequential o) => o -> ProcFun
parAwareProc opts = if (sequential opts) then procWithEmptyOutput else Turtle.procStrictWithErr

parAwareActions :: (HasSequential o, HasBatchSize o) => o -> [IO a] -> IO [a]
parAwareActions opts = if (sequential opts) then sequence else parBatchedActions (batchSize opts) []

parBatchedActions :: Int -> [a] -> [IO a] -> IO [a]
parBatchedActions _ done [] = return done
parBatchedActions batch done todo = do
  let doNow = take batch todo
  let remaining = drop batch todo
  doneNow <- (Turtle.single . shellToList . Turtle.parallel) doNow
  parBatchedActions batch (done ++ doneNow) remaining

inprocWithErrFun :: (T.Text -> IO ()) -> ProcInput -> Turtle.Shell Turtle.Line
inprocWithErrFun errFun (cmd, args, standardInput) = do
  result <- Turtle.inprocWithErr cmd args standardInput
  case result of
    Right ln -> return ln
    Left ln -> do
      (Turtle.liftIO . errFun . Turtle.lineToText) ln
      Turtle.empty

verboseTestFile :: (HasVerbosity o, HasBaseDir o) => o -> TChan LogMessage -> TurtlePath -> IO Bool
verboseTestFile opts ch p = do
  fileExists <- Turtle.testfile p
  let rel = relativeToBase opts p
  if fileExists
    then logVerbose opts ch $ Turtle.format ("Found '" % Turtle.fp % "'") rel
    else logVerbose opts ch $ Turtle.format ("Looked for but did not find '" % Turtle.fp % "'") rel
  return fileExists

needsRegeneration :: TurtlePath -> TurtlePath -> IO Bool
needsRegeneration src target = do
  targetExists <- Turtle.testfile target
  if not targetExists
    then return True
    else do
      -- Use try to handle race conditions where files may be modified/deleted
      -- between existence check and stat. Treat any IO error as "needs regeneration".
      result <- try $ do
        srcStat <- Turtle.stat src
        targetStat <- Turtle.stat target
        let srcMtime = Turtle.modificationTime srcStat
        let targetMtime = Turtle.modificationTime targetStat
        return (srcMtime > targetMtime)
      case result of
        Right needsRegen -> return needsRegen
        Left (_ :: IOException) -> return True

groupPairs' :: (Eq a, Ord a) => [(a, b)] -> [(a, [b])]
groupPairs' =
  mapMaybe
    ( \ll -> case ll of
        [] -> Nothing
        ((k, _) : _) -> Just (k, map snd ll)
    )
    . List.groupBy ((==) `on` fst)
    . List.sortBy (comparing fst)

groupPairs :: (Eq a, Ord a) => [(a, b)] -> Map.Map a [b]
groupPairs = Map.fromList . groupPairs'

pairBy :: (a -> b) -> [a] -> [(b, a)]
pairBy keyFun = map (\v -> (keyFun v, v))

groupValuesBy :: (Ord k, Ord v) => (v -> k) -> [v] -> Map.Map k [v]
groupValuesBy keyFun = groupPairs . pairBy keyFun

allYearsFileName :: TurtlePath
allYearsFileName = "all-years" <.> "journal"

directivesFile :: (HasBaseDir o) => o -> TurtlePath
directivesFile opts = turtleBaseDir opts </> "directives" <.> "journal"

lsDirs :: TurtlePath -> Turtle.Shell TurtlePath
lsDirs = onlyDirs . Turtle.ls

onlyDirs :: Turtle.Shell TurtlePath -> Turtle.Shell TurtlePath
onlyDirs = excludeHiddenFiles . excludeWeirdPaths . filterPathsByFileStatus Turtle.isDirectory

onlyFiles :: Turtle.Shell TurtlePath -> Turtle.Shell TurtlePath
onlyFiles = excludeHiddenFiles . filterPathsByFileStatus Turtle.isRegularFile

filterPathsByFileStatus :: (Turtle.FileStatus -> Bool) -> Turtle.Shell TurtlePath -> Turtle.Shell TurtlePath
filterPathsByFileStatus filepred files = do
  files' <- shellToList files
  filtered <- filterPathsByFileStatus' filepred [] files'
  Turtle.select filtered

filterPathsByFileStatus' :: (Turtle.FileStatus -> Bool) -> [TurtlePath] -> [TurtlePath] -> Turtle.Shell [TurtlePath]
filterPathsByFileStatus' _ acc [] = return acc
filterPathsByFileStatus' filepred acc (file : files) = do
  filestat <- Turtle.stat file
  let filtered = if (filepred filestat) then file : acc else acc
  filterPathsByFileStatus' filepred filtered files

filterPaths :: (TurtlePath -> IO Bool) -> [TurtlePath] -> Turtle.Shell [TurtlePath]
filterPaths = filterPaths' []

filterPaths' :: [TurtlePath] -> (TurtlePath -> IO Bool) -> [TurtlePath] -> Turtle.Shell [TurtlePath]
filterPaths' acc _ [] = return acc
filterPaths' acc filepred (file : files) = do
  shouldInclude <- Turtle.liftIO $ filepred file
  let filtered = if shouldInclude then file : acc else acc
  filterPaths' filtered filepred files

excludeHiddenFiles :: Turtle.Shell TurtlePath -> Turtle.Shell TurtlePath
excludeHiddenFiles paths = do
  p <- paths
  case (Turtle.match (Turtle.prefix ".") $ Turtle.format Turtle.fp $ Turtle.filename p) of
    [] -> Turtle.select [p]
    _ -> Turtle.select []

excludeWeirdPaths :: Turtle.Shell TurtlePath -> Turtle.Shell TurtlePath
excludeWeirdPaths = Turtle.findtree (Turtle.suffix $ Turtle.noneOf "_")

firstExistingFile :: [TurtlePath] -> IO (Maybe TurtlePath)
firstExistingFile files = do
  case files of
    [] -> return Nothing
    file : fs -> do
      exists <- Turtle.testfile file
      if exists then return (Just file) else firstExistingFile fs

basenameLine :: TurtlePath -> Turtle.Shell Turtle.Line
basenameLine path = case (Turtle.textToLine $ Turtle.format Turtle.fp $ Turtle.basename path) of
  Nothing -> Turtle.die $ Turtle.format ("Unable to determine basename from path: " % Turtle.fp % "\n") path
  Just bn -> return bn

buildFilename :: [Turtle.Line] -> T.Text -> TurtlePath
buildFilename identifiers ext = T.unpack (T.intercalate "-" (map Turtle.lineToText identifiers)) Turtle.<.> (T.unpack ext)

shellToList :: Turtle.Shell a -> Turtle.Shell [a]
shellToList files = Turtle.fold files Fold.list

writeFiles :: IO (Map.Map TurtlePath T.Text) -> IO [TurtlePath]
writeFiles fileMap = do
  m <- fileMap
  writeFiles' m

writeFiles' :: Map.Map TurtlePath T.Text -> IO [TurtlePath]
writeFiles' fileMap = do
  writeTextMap fileMap
  return $ Map.keys fileMap

writeTextMap :: Map.Map TurtlePath T.Text -> IO ()
writeTextMap = Map.foldlWithKey (\a k v -> a <> T.writeFile k v) (return ())

changeExtension :: T.Text -> TurtlePath -> TurtlePath
changeExtension ext path = (Turtle.dropExtension path) Turtle.<.> (T.unpack ext)

changePathAndExtension :: TurtlePath -> T.Text -> TurtlePath -> TurtlePath
changePathAndExtension newOutputLocation newExt = (changeOutputPath newOutputLocation) . (changeExtension newExt)

changeOutputPath :: TurtlePath -> TurtlePath -> TurtlePath
changeOutputPath newOutputLocation srcFile = mconcat $ map changeSrcDir $ Turtle.splitDirectories srcFile
  where
    changeSrcDir file = if file == "1-in/" || file == "2-preprocessed/" then newOutputLocation else file

listOwners :: (HasBaseDir o) => o -> Turtle.Shell TurtlePath
listOwners opts = fmap Turtle.basename $ lsDirs $ (turtleBaseDir opts) </> "import"

intPath :: Integer -> TurtlePath
intPath = T.unpack . (Turtle.format Turtle.d)

includeYears :: TChan LogMessage -> TurtlePath -> IO [Integer]
includeYears ch includeFile = do
  txt <- T.readFile includeFile
  case includeYears' txt of
    Left msg -> do
      channelErrLn ch msg
      return []
    Right years -> return years

includeYears' :: T.Text -> Either T.Text [Integer]
includeYears' txt = case partitionEithers (includeYears'' txt) of
  (errors, []) -> do
    let msg = Turtle.format ("Unable to extract years from the following text:\n" % Turtle.s % "\nErrors:\n" % Turtle.s) txt (T.intercalate "\n" $ map T.pack errors)
    Left msg
  (_, years) -> Right years

includeYears'' :: T.Text -> [Either String Integer]
includeYears'' txt = map extractDigits (T.lines txt)

extractDigits :: T.Text -> Either String Integer
extractDigits txt = fmap fst $ (T.decimal . (T.filter isDigit)) txt