pwncash/lib/Internal/History.hs

module Internal.History
  ( readHistStmt
  , readHistTransfer
  , insertHistory
  , splitHistory
  )
where

import Control.Monad.Except
import Data.Csv
import Data.Foldable
import Database.Persist.Monad hiding (get)
import Internal.Database
import Internal.Types.Main
import Internal.Utils
import RIO hiding (to)
import qualified RIO.ByteString.Lazy as BL
import RIO.FilePath
import qualified RIO.List as L
import qualified RIO.Map as M
import qualified RIO.NonEmpty as NE
import RIO.State
import qualified RIO.Text as T
import RIO.Time
import qualified RIO.Vector as V

-- readHistory
--   :: (MonadInsertError m, MonadFinance m, MonadUnliftIO m)
--   => FilePath
--   -> [History]
--   -> m [(CommitR, [DeferredTx])]
-- readHistory root hs = do
--   let (ts, ss) = splitHistory hs
--   ts' <- catMaybes <$> mapErrorsIO readHistTransfer ts
--   ss' <- catMaybes <$> mapErrorsIO (readHistStmt root) ss
--   return $ ts' ++ ss'

readHistTransfer
  :: (MonadInsertError m, MonadFinance m)
  => HistTransfer
  -> m (Maybe (CommitR, [DeferredTx]))
readHistTransfer
  m@Transfer
    { transFrom = from
    , transTo = to
    , transCurrency = u
    , transAmounts = amts
    } = do
    whenHash_ CTManual m $ do
      bounds <- askDBState kmStatementInterval
      let precRes = lookupCurrencyPrec u
      let go Amount {amtWhen, amtValue, amtDesc} = do
            let dayRes = liftExcept $ expandDatePat bounds amtWhen
            (days, precision) <- combineError dayRes precRes (,)
            let tx day = txPair day from to u (roundPrecision precision amtValue) amtDesc
            return $ fmap tx days
      concat <$> mapErrors go amts

groupKey :: Ord c => (a -> c) -> [(a, b)] -> [(a, [b])]
groupKey f = fmap go . NE.groupAllWith (f . fst)
  where
    go xs@((c, _) :| _) = (c, NE.toList $ fmap snd xs)

readHistStmt
  :: (MonadUnliftIO m, MonadFinance m)
  => FilePath
  -> Statement
  -> m (Maybe (CommitR, [DeferredTx]))
readHistStmt root i = whenHash_ CTImport i $ do
  bs <- readImport root i
  bounds <- askDBState kmStatementInterval
  return $ filter (inDaySpan bounds . dtxDate) bs

splitHistory :: [History] -> ([HistTransfer], [Statement])
splitHistory = partitionEithers . fmap go
  where
    go (HistTransfer x) = Left x
    go (HistStatement x) = Right x

insertHistory
  :: (MonadInsertError m, MonadSqlQuery m, MonadFinance m)
  => [(CommitR, [DeferredTx])]
  -> m ()
insertHistory hs = do
  bs <- balanceTxs $ concatMap (\(c, xs) -> fmap (c,) xs) hs
  forM_ (groupKey (\(CommitR h _) -> h) bs) $ \(c, ts) -> do
    ck <- insert c
    mapM_ (insertTx ck) ts

--------------------------------------------------------------------------------
-- low-level transaction stuff

-- TODO tags here?
txPair
  :: Day
  -> AcntID
  -> AcntID
  -> CurID
  -> Rational
  -> T.Text
  -> DeferredTx
txPair day from to cur val desc =
  Tx
    { dtxDescr = desc
    , dtxDate = day
    , dtxEntries =
        [ EntrySet
            { desTotalValue = val
            , desCurrency = cur
            , desFromEntry0 = entry from
            , desToEntryBal = entry to
            , desFromEntries = []
            , desToEntries = []
            }
        ]
    }
  where
    entry a =
      Entry
        { eAcnt = a
        , eValue = ()
        , eComment = ""
        , eTags = []
        }

resolveTx :: (MonadInsertError m, MonadFinance m) => BalTx -> m KeyTx
resolveTx t@Tx {dtxEntries = ss} =
  (\kss -> t {dtxEntries = kss}) <$> mapErrors resolveEntry ss

insertTx :: MonadSqlQuery m => CommitRId -> KeyTx -> m ()
insertTx c Tx {dtxDate = d, dtxDescr = e, dtxEntries = ss} = do
  k <- insert $ TransactionR c d e
  mapM_ (insertEntry k) ss

--------------------------------------------------------------------------------
-- Statements

-- TODO this probably won't scale well (pipes?)
readImport :: (MonadUnliftIO m, MonadFinance m) => FilePath -> Statement -> m [DeferredTx]
readImport root Statement {stmtTxOpts, stmtParsers, stmtSkipLines, stmtDelim, stmtPaths} = do
  let ores = compileOptions stmtTxOpts
  let cres = combineErrors $ compileMatch <$> stmtParsers
  (compiledOptions, compiledMatches) <- liftIOExcept $ combineError ores cres (,)
  let readStmt = readImport_ stmtSkipLines stmtDelim compiledOptions
  records <- L.sort . concat <$> mapErrorsIO readStmt paths
  m <- askDBState kmCurrency
  fromEither $
    flip runReader m $
      runExceptT $
        matchRecords compiledMatches records
  where
    paths = (root </>) <$> stmtPaths

readImport_
  :: MonadUnliftIO m
  => Natural
  -> Word
  -> TxOptsRe
  -> FilePath
  -> m [TxRecord]
readImport_ n delim tns p = do
  res <- tryIO $ BL.readFile p
  bs <- fromEither $ first (InsertException . (: []) . InsertIOError . showT) res
  case decodeByNameWithP (parseTxRecord p tns) opts $ skip bs of
    Left m -> throwIO $ InsertException [ParseError $ T.pack m]
    Right (_, v) -> return $ catMaybes $ V.toList v
  where
    opts = defaultDecodeOptions {decDelimiter = fromIntegral delim}
    skip = BL.intercalate "\n" . L.drop (fromIntegral n) . BL.split 10

-- TODO handle this better, this maybe thing is a hack to skip lines with
-- blank dates but will likely want to make this more flexible
parseTxRecord :: FilePath -> TxOptsRe -> NamedRecord -> Parser (Maybe TxRecord)
parseTxRecord p TxOpts {toDate, toAmountFmt, toDesc, toAmount, toOther, toDateFmt} r = do
  d <- r .: T.encodeUtf8 toDate
  if d == ""
    then return Nothing
    else do
      a <- parseRational toAmountFmt =<< r .: T.encodeUtf8 toAmount
      e <- r .: T.encodeUtf8 toDesc
      os <- M.fromList <$> mapM (\n -> (n,) <$> r .: T.encodeUtf8 n) toOther
      d' <- parseTimeM True defaultTimeLocale (T.unpack toDateFmt) d
      return $ Just $ TxRecord d' a e os p

-- TODO need to somehow balance temporally here (like I do in the budget for
-- directives that "pay off" a balance)
matchRecords :: [MatchRe] -> [TxRecord] -> InsertExceptT CurrencyM [DeferredTx]
matchRecords ms rs = do
  (matched, unmatched, notfound) <- matchAll (matchPriorities ms) rs
  case (matched, unmatched, notfound) of
    (ms_, [], []) -> return ms_ -- liftInner $ combineErrors $ fmap balanceTx ms_
    (_, us, ns) -> throwError $ InsertException [StatementError us ns]

matchPriorities :: [MatchRe] -> [MatchGroup]
matchPriorities =
  fmap matchToGroup
    . L.groupBy (\a b -> spPriority a == spPriority b)
    . L.sortOn (Down . spPriority)

matchToGroup :: [MatchRe] -> MatchGroup
matchToGroup ms =
  uncurry MatchGroup $
    first (L.sortOn spDate) $
      L.partition (isJust . spDate) ms

data MatchGroup = MatchGroup
  { mgDate :: ![MatchRe]
  , mgNoDate :: ![MatchRe]
  }
  deriving (Show)

data Zipped a = Zipped ![a] ![a]

data Unzipped a = Unzipped ![a] ![a] ![a]

initZipper :: [a] -> Zipped a
initZipper = Zipped []

resetZipper :: Zipped a -> Zipped a
resetZipper = initZipper . recoverZipper

recoverZipper :: Zipped a -> [a]
recoverZipper (Zipped as bs) = reverse as ++ bs

zipperSlice
  :: (a -> b -> Ordering)
  -> b
  -> Zipped a
  -> Either (Zipped a) (Unzipped a)
zipperSlice f x = go
  where
    go z@(Zipped _ []) = Left z
    go z@(Zipped bs (a : as)) =
      case f a x of
        GT -> go $ Zipped (a : bs) as
        EQ -> Right $ goEq (Unzipped bs [a] as)
        LT -> Left z
    goEq z@(Unzipped _ _ []) = z
    goEq z@(Unzipped bs cs (a : as)) =
      case f a x of
        GT -> goEq $ Unzipped (a : bs) cs as
        EQ -> goEq $ Unzipped bs (a : cs) as
        LT -> z

zipperMatch
  :: Unzipped MatchRe
  -> TxRecord
  -> InsertExceptT CurrencyM (Zipped MatchRe, MatchRes DeferredTx)
zipperMatch (Unzipped bs cs as) x = go [] cs
  where
    go _ [] = return (Zipped bs $ cs ++ as, MatchFail)
    go prev (m : ms) = do
      res <- matches m x
      case res of
        MatchFail -> go (m : prev) ms
        skipOrPass ->
          let ps = reverse prev
              ms' = maybe ms (: ms) (matchDec m)
           in return (Zipped bs $ ps ++ ms' ++ as, skipOrPass)

zipperMatch'
  :: Zipped MatchRe
  -> TxRecord
  -> InsertExceptT CurrencyM (Zipped MatchRe, MatchRes DeferredTx)
zipperMatch' z x = go z
  where
    go (Zipped bs (a : as)) = do
      res <- matches a x
      case res of
        MatchFail -> go (Zipped (a : bs) as)
        skipOrPass ->
          return (Zipped (maybe bs (: bs) $ matchDec a) as, skipOrPass)
    go z' = return (z', MatchFail)

matchDec :: MatchRe -> Maybe MatchRe
matchDec m = case spTimes m of
  Just 1 -> Nothing
  Just n -> Just $ m {spTimes = Just $ n - 1}
  Nothing -> Just m

matchAll :: [MatchGroup] -> [TxRecord] -> InsertExceptT CurrencyM ([DeferredTx], [TxRecord], [MatchRe])
matchAll = go ([], [])
  where
    go (matched, unused) gs rs = case (gs, rs) of
      (_, []) -> return (matched, [], unused)
      ([], _) -> return (matched, rs, unused)
      (g : gs', _) -> do
        (ts, unmatched, us) <- matchGroup g rs
        go (ts ++ matched, us ++ unused) gs' unmatched

matchGroup :: MatchGroup -> [TxRecord] -> InsertExceptT CurrencyM ([DeferredTx], [TxRecord], [MatchRe])
matchGroup MatchGroup {mgDate = ds, mgNoDate = ns} rs = do
  (md, rest, ud) <- matchDates ds rs
  (mn, unmatched, un) <- matchNonDates ns rest
  return (md ++ mn, unmatched, filter ((/= Nothing) . spTimes) $ ud ++ un)

matchDates :: [MatchRe] -> [TxRecord] -> InsertExceptT CurrencyM ([DeferredTx], [TxRecord], [MatchRe])
matchDates ms = go ([], [], initZipper ms)
  where
    go (matched, unmatched, z) [] =
      return
        ( catMaybes matched
        , reverse unmatched
        , recoverZipper z
        )
    go (matched, unmatched, z) (r : rs) =
      case zipperSlice findDate r z of
        Left zipped -> go (matched, r : unmatched, zipped) rs
        Right unzipped -> do
          (z', res) <- zipperMatch unzipped r
          let (m, u) = case res of
                (MatchPass p) -> (Just p : matched, unmatched)
                MatchSkip -> (Nothing : matched, unmatched)
                MatchFail -> (matched, r : unmatched)
          go (m, u, z') rs
    findDate m r = maybe EQ (`compareDate` trDate r) $ spDate m

matchNonDates :: [MatchRe] -> [TxRecord] -> InsertExceptT CurrencyM ([DeferredTx], [TxRecord], [MatchRe])
matchNonDates ms = go ([], [], initZipper ms)
  where
    go (matched, unmatched, z) [] =
      return
        ( catMaybes matched
        , reverse unmatched
        , recoverZipper z
        )
    go (matched, unmatched, z) (r : rs) = do
      (z', res) <- zipperMatch' z r
      let (m, u) = case res of
            MatchPass p -> (Just p : matched, unmatched)
            MatchSkip -> (Nothing : matched, unmatched)
            MatchFail -> (matched, r : unmatched)
       in go (m, u, resetZipper z') rs

-- TDOO should use a better type here to squish down all the entry sets
-- which at this point in the chain should not be necessary
balanceTxs
  :: (MonadInsertError m, MonadFinance m)
  => [(CommitR, DeferredTx)]
  -> m [(CommitR, KeyTx)]
balanceTxs ts = do
  keyts <- mapErrors resolveTx balTs
  return $ zip cs keyts
  where
    (cs, ts') = L.unzip $ L.sortOn (dtxDate . snd) ts
    go bals t@Tx {dtxEntries} =
      second (\es -> t {dtxEntries = concat es}) $
        L.mapAccumL balanceEntrySet bals dtxEntries
    balTs = snd $ L.mapAccumL go M.empty ts'

type EntryBals = M.Map (AcntID, CurID) Rational

-- TODO might be faster to also do all the key stuff here since currency
-- will be looked up for every entry rather then the entire entry set
balanceEntrySet :: EntryBals -> DeferredEntrySet -> (EntryBals, [BalEntry])
balanceEntrySet
  bals
  EntrySet
    { desFromEntry0
    , desFromEntries
    , desToEntryBal
    , desToEntries
    , desCurrency
    , desTotalValue
    } = flipTup $ runState doBalAll bals
    where
      flipTup (a, b) = (b, a)
      doEntries es tot e0 = do
        es' <- state (\b -> flipTup $ L.mapAccumL (balanceEntry desCurrency) b es)
        let val0 = tot - entrySum es'
        modify $ mapAdd_ (eAcnt e0, desCurrency) val0
        return $ e0 {eValue = val0} : es'
      doBalAll = do
        fes <- doEntries desFromEntries desTotalValue desFromEntry0
        tes <- doEntries desToEntries (-desTotalValue) desToEntryBal
        return $ toFull <$> fes ++ tes
      toFull e = FullEntry {feEntry = e, feCurrency = desCurrency}

entrySum :: Num v => [Entry a v t] -> v
entrySum = sum . fmap eValue

balanceEntry
  :: CurID
  -> EntryBals
  -> Entry AcntID (Deferred Rational) TagID
  -> (EntryBals, Entry AcntID Rational TagID)
balanceEntry curID bals e@Entry {eValue = Deferred toBal v, eAcnt}
  | toBal = (bals, e {eValue = v})
  | otherwise = (bals', e {eValue = newVal})
  where
    key = (eAcnt, curID)
    curBal = M.findWithDefault 0 key bals
    newVal = v - curBal
    bals' = mapAdd_ key newVal bals

-- -- reimplementation from future version :/
-- mapAccumM
--   :: Monad m
--   => (s -> a -> m (s, b))
--   -> s
--   -> [a]
--   -> m (s, [b])
-- mapAccumM f s xs = foldrM go (s, []) xs
--   where
--     go x (s', acc) = second (: acc) <$> f s' x