ImportHelpersTurtle.hs

{-# LANGUAGE OverloadedStrings #-}

-- | Functions which currently uses TurtlePath and will be replaced with Path eventually
module Hledger.Flow.Import.ImportHelpersTurtle (
    allYearIncludeFiles
  , extractImportDirs
  , extraIncludesForFile
  , groupIncludeFiles
  , groupAndWriteIncludeFiles
  , includePreamble
  , toIncludeFiles
  , toIncludeLine
  , writeIncludesUpTo
  , writeToplevelAllYearsInclude
  , yearsIncludeMap
  )
 where

import Hledger.Flow.PathHelpers (TurtlePath)
import Hledger.Flow.DocHelpers (docURL)
import Hledger.Flow.Common (allYearsFileName, filterPaths, groupValuesBy, writeFiles, writeFiles')
import Hledger.Flow.BaseDir (relativeToBase, relativeToBase', turtleBaseDir)
import Hledger.Flow.Logging (logVerbose)

import Hledger.Flow.Types
import Hledger.Flow.Import.Types

import qualified Data.List as List (nub, sort)
import qualified Data.Text as T
import qualified Data.Map.Strict as Map

import qualified Turtle
import Turtle ((%), (</>), (<.>))

import Control.Concurrent.STM (TChan)
import Data.Maybe (fromMaybe)

extractImportDirs :: TurtlePath -> Either T.Text ImportDirs
extractImportDirs inputFile = do
  case importDirBreakdown inputFile of
    [bd,owner,bank,account,filestate,year] -> Right $ ImportDirs bd owner bank account filestate year
    _ -> do
      Left $ Turtle.format ("I couldn't find the right number of directories between \"import\" and the input file:\n"%Turtle.fp
                      %"\n\nhledger-flow expects to find input files in this structure:\n"%
                      "import/owner/bank/account/filestate/year/trxfile\n\n"%
                      "Have a look at the documentation for a detailed explanation:\n"%Turtle.s) inputFile (docURL "input-files")

importDirBreakdown ::  TurtlePath -> [TurtlePath]
importDirBreakdown = importDirBreakdown' []

importDirBreakdown' :: [TurtlePath] -> TurtlePath -> [TurtlePath]
importDirBreakdown' acc path = do
  let dir = Turtle.directory path
  if Turtle.dirname dir == "import" || (Turtle.dirname dir == "")
    then dir:acc
    else importDirBreakdown' (dir:acc) $ Turtle.parent dir

groupIncludeFiles :: [TurtlePath] -> (TurtleFileBundle, TurtleFileBundle)
groupIncludeFiles = allYearIncludeFiles . groupIncludeFilesPerYear . filter isJournalFile

isJournalFile :: TurtlePath -> Bool
isJournalFile f = Turtle.extension f == Just "journal"

allYearIncludeFiles :: TurtleFileBundle -> (TurtleFileBundle, TurtleFileBundle)
allYearIncludeFiles m = (m, yearsIncludeMap $ Map.keys m)

yearsIncludeMap :: [TurtlePath] -> TurtleFileBundle
yearsIncludeMap = groupValuesBy allYearsPath

allYearsPath :: TurtlePath -> TurtlePath
allYearsPath = allYearsPath' Turtle.directory

allYearsPath' :: (TurtlePath -> TurtlePath) -> TurtlePath -> TurtlePath
allYearsPath' dir p = dir p </> allYearsFileName

includeFileName :: TurtlePath -> TurtlePath
includeFileName = (<.> "journal") . Turtle.fromText . (Turtle.format (Turtle.fp%"-include")) . Turtle.dirname

groupIncludeFilesPerYear :: [TurtlePath] -> TurtleFileBundle
groupIncludeFilesPerYear [] = Map.empty
groupIncludeFilesPerYear ps@(p:_) = case extractImportDirs p of
    Right _ -> groupValuesBy initialIncludeFilePath ps
    Left  _ -> groupValuesBy parentIncludeFilePath ps

initialIncludeFilePath :: TurtlePath -> TurtlePath
initialIncludeFilePath p = (Turtle.parent . Turtle.parent . Turtle.parent) p </> includeFileName p

parentIncludeFilePath :: TurtlePath -> TurtlePath
parentIncludeFilePath p = (Turtle.parent . Turtle.parent) p </> Turtle.filename p

toIncludeFiles :: (HasBaseDir o, HasVerbosity o) => o -> TChan LogMessage -> TurtleFileBundle -> IO (Map.Map TurtlePath T.Text)
toIncludeFiles opts ch m = do
  preMap  <- extraIncludes opts ch (Map.keys m) ["opening.journal"] ["pre-import.journal"] []
  postMap <- extraIncludes opts ch (Map.keys m) ["closing.journal"] ["post-import.journal"] ["prices.journal"]
  return $ (addPreamble . toIncludeFiles' preMap postMap) m

toIncludeFiles' :: TurtleFileBundle -> TurtleFileBundle -> TurtleFileBundle -> Map.Map TurtlePath T.Text
toIncludeFiles' preMap postMap = Map.mapWithKey $ generatedIncludeText preMap postMap

addPreamble :: Map.Map TurtlePath T.Text -> Map.Map TurtlePath T.Text
addPreamble = Map.map (\txt -> includePreamble <> "\n" <> txt)

toIncludeLine :: TurtlePath -> TurtlePath -> T.Text
toIncludeLine base file = Turtle.format ("!include "%Turtle.fp) $ relativeToBase' base file

generatedIncludeText :: TurtleFileBundle -> TurtleFileBundle -> TurtlePath -> [TurtlePath] -> T.Text
generatedIncludeText preMap postMap outputFile fs = do
  let preFiles = fromMaybe [] $ Map.lookup outputFile preMap
  let files = List.nub . List.sort $ fs
  let postFiles = fromMaybe [] $ Map.lookup outputFile postMap
  let lns = map (toIncludeLine $ Turtle.directory outputFile) $ preFiles ++ files ++ postFiles
  T.intercalate "\n" $ lns ++ [""]

includePreamble :: T.Text
includePreamble = "### Generated by hledger-flow - DO NOT EDIT ###\n"

groupAndWriteIncludeFiles :: (HasBaseDir o, HasVerbosity o) => o -> TChan LogMessage -> [TurtlePath] -> IO [TurtlePath]
groupAndWriteIncludeFiles opts ch = writeFileMap opts ch . groupIncludeFiles

writeFileMap :: (HasBaseDir o, HasVerbosity o) => o -> TChan LogMessage -> (TurtleFileBundle, TurtleFileBundle) -> IO [TurtlePath]
writeFileMap opts ch (m, allYears) = do
  _ <- writeFiles' $ (addPreamble . toIncludeFiles' Map.empty Map.empty) allYears
  writeFiles . (toIncludeFiles opts ch) $ m

writeIncludesUpTo :: (HasBaseDir o, HasVerbosity o) => o -> TChan LogMessage -> TurtlePath -> [TurtlePath] -> IO [TurtlePath]
writeIncludesUpTo _ _ _ [] = return []
writeIncludesUpTo opts ch stopAt journalFiles = do
  let shouldStop = any (\dir -> dir == stopAt) $ map Turtle.parent journalFiles
  if shouldStop
    then return journalFiles
    else do
      newJournalFiles <- groupAndWriteIncludeFiles opts ch journalFiles
      writeIncludesUpTo opts ch stopAt newJournalFiles

writeToplevelAllYearsInclude :: (HasBaseDir o, HasVerbosity o) => o -> IO [TurtlePath]
writeToplevelAllYearsInclude opts = do
  let directivesFile = turtleBaseDir opts </> "directives" <.> "journal"
  directivesExists <- Turtle.testfile directivesFile
  let preMap = if directivesExists then Map.singleton (turtleBaseDir opts </> allYearsFileName) [directivesFile] else Map.empty
  let allTop = Map.singleton (turtleBaseDir opts </> allYearsFileName) ["import" </> allYearsFileName]
  writeFiles' $ (addPreamble . toIncludeFiles' preMap Map.empty) allTop

extraIncludes :: (HasBaseDir o, HasVerbosity o) => o -> TChan LogMessage -> [TurtlePath] -> [T.Text] -> [TurtlePath] -> [TurtlePath] -> IO TurtleFileBundle
extraIncludes opts ch = extraIncludes' opts ch Map.empty

extraIncludes' :: (HasBaseDir o, HasVerbosity o) => o -> TChan LogMessage -> TurtleFileBundle -> [TurtlePath] -> [T.Text] -> [TurtlePath] -> [TurtlePath] -> IO TurtleFileBundle
extraIncludes' _ _ acc [] _ _ _ = return acc
extraIncludes' opts ch acc (file:files) extraSuffixes manualFiles prices = do
  extra <- extraIncludesForFile opts ch file extraSuffixes manualFiles prices
  extraIncludes' opts ch (Map.unionWith (++) acc extra) files extraSuffixes manualFiles prices

extraIncludesForFile :: (HasVerbosity o, HasBaseDir o) => o -> TChan LogMessage -> TurtlePath -> [T.Text] -> [TurtlePath] -> [TurtlePath] -> IO TurtleFileBundle
extraIncludesForFile opts ch file extraSuffixes manualFiles prices = do
  let dirprefix = Turtle.fromText $ fst $ T.breakOn "-" $ Turtle.format Turtle.fp $ Turtle.basename file
  let fileNames = map (Turtle.fromText . Turtle.format (Turtle.fp % "-" % Turtle.s) dirprefix) extraSuffixes
  let suffixFiles = map (Turtle.directory file </>) fileNames
  let suffixDirFiles = map (Turtle.directory file </> "_manual_" </> dirprefix </>) manualFiles
  let priceFiles = map (Turtle.directory file </> ".." </> "prices" </> dirprefix </>) prices
  let extraFiles = suffixFiles ++ suffixDirFiles ++ priceFiles
  filtered <- Turtle.single $ filterPaths Turtle.testfile extraFiles
  let logMsg = Turtle.format ("Looking for possible extra include files for '"%Turtle.fp%"' among these "%Turtle.d%" options: "%Turtle.s%". Found "%Turtle.d%": "%Turtle.s)
               (relativeToBase opts file) (length extraFiles) (Turtle.repr $ relativeFilesAsText opts extraFiles)
               (length filtered) (Turtle.repr $ relativeFilesAsText opts filtered)
  logVerbose opts ch logMsg
  return $ Map.fromList [(file, filtered)]

relativeFilesAsText :: HasBaseDir o => o -> [TurtlePath] -> [T.Text]
relativeFilesAsText opts = map (Turtle.format Turtle.fp . relativeToBase opts)