ImportHelpersTurtle.hs

{-# LANGUAGE OverloadedStrings #-}

-- | Generate journal include files and their relative include directives.
module Hledger.Flow.Import.ImportHelpersTurtle
  ( allYearIncludeFiles,
    extractImportDirs,
    extraIncludesForFile,
    groupIncludeFiles,
    groupAndWriteIncludeFiles,
    includePreamble,
    toIncludeFiles,
    toIncludeLine,
    writeIncludesUpTo,
    writeToplevelAllYearsInclude,
    yearsIncludeMap,
  )
where

import Control.Concurrent.STM (TChan)
import Control.Monad (unless)
import Data.Either (isRight)
import qualified Data.List as List (isPrefixOf, nub, sort)
import qualified Data.Map.Strict as Map
import Data.Maybe (fromMaybe)
import qualified Data.Text as T
import Hledger.Flow.BaseDir (relativeToBase, relativeToBase', turtleBaseDir)
import Hledger.Flow.Common (allYearsFileName, directivesFile, filterPaths, groupValuesBy, writeFiles, writeFiles')
import Hledger.Flow.Import.ImportHelpers (extractImportDirs, groupIncludesUpTo)
import Hledger.Flow.Import.Types
import Hledger.Flow.Logging (logVerbose)
import Hledger.Flow.PathHelpers (AbsDir, AbsFile)
import Hledger.Flow.Types
import qualified Path
import qualified Path.IO as PathIO
import System.Directory (doesFileExist)
import qualified System.FilePath as FilePath
import System.FilePath ((</>))

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

isJournalFile :: FilePath -> Bool
isJournalFile f = FilePath.takeExtension f == ".journal"

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

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

allYearsPath :: FilePath -> FilePath
allYearsPath journal = fst (FilePath.splitFileName journal) </> allYearsFileName

includeFileName :: FilePath -> FilePath
includeFileName journal = FilePath.takeFileName (FilePath.takeDirectory journal) ++ "-include.journal"

groupIncludeFilesPerYear :: [FilePath] -> TurtleFileBundle
groupIncludeFilesPerYear [] = Map.empty
groupIncludeFilesPerYear ps@(p : _)
  | maybe False (Path.prjSomeBase (isRight . extractImportDirs)) (Path.parseSomeFile p) =
      groupValuesBy initialIncludeFilePath ps
  | otherwise = groupValuesBy parentIncludeFilePath ps

initialIncludeFilePath :: FilePath -> FilePath
initialIncludeFilePath p = ancestorDir 3 p </> includeFileName p

parentIncludeFilePath :: FilePath -> FilePath
parentIncludeFilePath p = ancestorDir 2 p </> FilePath.takeFileName p

-- Keep the original path components (including separators) when constructing
-- include-map keys. Relative keys have historically been rooted at "./".
ancestorDir :: Int -> FilePath -> FilePath
ancestorDir levels file = prefix </> FilePath.joinPath parents
  where
    (drive, rest) = FilePath.splitDrive file
    parts = FilePath.splitPath rest
    parents = take (length parts - levels) parts
    prefix = case parents of
      first : _ | first == "./" || first == "../" -> drive
      _ | null drive -> "./"
      _ -> drive

toIncludeFiles :: (HasBaseDir o, HasVerbosity o) => o -> TChan LogMessage -> TurtleFileBundle -> IO (Map.Map FilePath 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 FilePath T.Text
toIncludeFiles' preMap postMap = Map.mapWithKey $ generatedIncludeText preMap postMap

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

toIncludeLine :: FilePath -> FilePath -> T.Text
toIncludeLine base file = "include " <> T.pack (map forwardSlash (relativeToBase' base file))
  where
    forwardSlash c = if FilePath.isPathSeparator c then '/' else c

generatedIncludeText :: TurtleFileBundle -> TurtleFileBundle -> FilePath -> [FilePath] -> 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 $ FilePath.takeDirectory 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 -> [FilePath] -> IO [FilePath]
groupAndWriteIncludeFiles opts ch = writeFileMap opts ch . groupIncludeFiles

writeFileMap :: (HasBaseDir o, HasVerbosity o) => o -> TChan LogMessage -> (TurtleFileBundle, TurtleFileBundle) -> IO [FilePath]
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 -> AbsDir -> [AbsFile] -> IO [AbsFile]
writeIncludesUpTo opts ch stopAt journals = do
  let base = baseDir opts
      journalFiles = filter (isJournalFile . Path.toFilePath) journals
  boundary <- Path.stripProperPrefix base stopAt
  unless (all (Path.isProperPrefixOf stopAt) journalFiles) $
    ioError (userError "Journal file outside the requested include directory")
  inputs <- mapM (Path.stripProperPrefix base) journalFiles
  let groups = groupIncludesUpTo boundary inputs
      absolute file = Path.toFilePath (base Path.</> file)
      includes = Map.mapKeys absolute (Map.map (map absolute) groups)
      years = yearsIncludeMap (Map.keys includes)
      allAtStop = Path.toFilePath stopAt </> allYearsFileName
  _ <- writeFileMap opts ch (includes, Map.insertWith (++) allAtStop [] years)
  return [base Path.</> f | f <- Map.keys groups, Path.parent f == boundary]

writeToplevelAllYearsInclude :: (HasBaseDir o, HasVerbosity o) => o -> IO [FilePath]
writeToplevelAllYearsInclude opts = do
  let directives = directivesFile opts
  directivesExists <- PathIO.doesFileExist directives
  let preMap = if directivesExists then Map.singleton (turtleBaseDir opts </> allYearsFileName) [Path.toFilePath directives] 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 -> [FilePath] -> [T.Text] -> [FilePath] -> [FilePath] -> IO TurtleFileBundle
extraIncludes opts ch = extraIncludes' opts ch Map.empty

extraIncludes' :: (HasBaseDir o, HasVerbosity o) => o -> TChan LogMessage -> TurtleFileBundle -> [FilePath] -> [T.Text] -> [FilePath] -> [FilePath] -> 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 -> FilePath -> [T.Text] -> [FilePath] -> [FilePath] -> IO TurtleFileBundle
extraIncludesForFile opts ch file extraSuffixes manualFiles prices = do
  let name = FilePath.takeFileName file
      stem = if "." `List.isPrefixOf` name then name else FilePath.dropExtensions name
      dirprefix = takeWhile (/= '-') stem
      fileNames = [dirprefix ++ "-" ++ T.unpack suffix | suffix <- extraSuffixes]
  let suffixFiles = map (FilePath.takeDirectory file </>) fileNames
  let suffixDirFiles = map (((FilePath.takeDirectory file </> "_manual_") </> dirprefix) </>) manualFiles
  let priceFiles = map ((((FilePath.takeDirectory file </> "..") </> "prices") </> dirprefix) </>) prices
  let extraFiles = suffixFiles ++ suffixDirFiles ++ priceFiles
  filtered <- filterPaths doesFileExist extraFiles
  let logMsg =
        "Looking for possible extra include files for '" <> T.pack (relativeToBase opts file)
          <> "' among these " <> T.pack (show (length extraFiles)) <> " options: "
          <> T.pack (show (relativeFilesAsText opts extraFiles))
          <> ". Found " <> T.pack (show (length filtered)) <> ": "
          <> T.pack (show (relativeFilesAsText opts filtered))
  logVerbose opts ch logMsg
  return $ Map.fromList [(file, filtered)]

relativeFilesAsText :: (HasBaseDir o) => o -> [FilePath] -> [T.Text]
relativeFilesAsText opts = map (T.pack . relativeToBase opts)