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 [DeferredTx CommitR]
readHistTransfer
  m@Transfer
    { transFrom = from
    , transTo = to
    , transCurrency = u
    , transAmounts = amts
    } =
    whenHash0 CTManual m [] $ \c -> 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 c 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)

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

readHistStmt
  :: (MonadUnliftIO m, MonadFinance m)
  => FilePath
  -> Statement
  -> m [DeferredTx CommitR]
readHistStmt root i = whenHash0 CTImport i [] $ \c -> do
  bs <- readImport root i
  bounds <- askDBState kmStatementInterval
  return $ filter (inDaySpan bounds . txDate) $ fmap (\t -> t {txCommit = c}) 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)
  => [DeferredTx CommitR]
  -> m ()
insertHistory hs = do
  bs <- balanceTxs hs
  forM_ (groupWith txCommit bs) $ \(c, ts) -> do
    ck <- insert c
    mapM_ (insertTx ck) ts

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

-- TODO tags here?
txPair
  :: CommitR
  -> Day
  -> AcntID
  -> AcntID
  -> CurID
  -> Rational
  -> T.Text
  -> DeferredTx CommitR
txPair commit day from to cur val desc =
  Tx
    { txDescr = desc
    , txDate = day
    , txCommit = commit
    , txEntries =
        [ EntrySet
            { esTotalValue = -val
            , esCurrency = cur
            , esFrom = HalfEntrySet {hesPrimary = entry from, hesOther = []}
            , esTo = HalfEntrySet {hesPrimary = entry to, hesOther = []}
            }
        ]
    }
  where
    entry a =
      Entry
        { eAcnt = a
        , eValue = ()
        , eComment = ""
        , eTags = []
        }

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

insertTx :: MonadSqlQuery m => CommitRId -> (KeyTx CommitR) -> m ()
insertTx c Tx {txDate = d, txDescr = e, txEntries = ss} = do
  let anyDeferred = any (isJust . feDeferred) ss
  k <- insert $ TransactionR c d e anyDeferred
  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

balanceTxs
  :: (MonadInsertError m, MonadFinance m)
  => [EntryBin]
  -> m ([UpdateEntry EntryRId Rational], [KeyTx CommitR])
balanceTxs es =
  (first concat . partitionEithers . catMaybes)
    <$> evalStateT (mapM go $ L.sortOn binDate es) M.empty
  where
    go (ToUpdate utx) = (Just . Left) <$> rebalanceEntrySet utx
    go (ToRead ReadEntry {reCurrency, reAcnt, reValue}) = do
      modify $ mapAdd_ (reAcnt, reCurrency) reValue
      return Nothing
    go (ToInsert (t@Tx {txEntries, txDate})) =
      (\es -> Just $ Right $ t {txEntries = concat es})
        <$> mapM (balanceEntrySet txDate) txEntries

binDate :: EntryBin -> Day
binDate (ToUpdate (UpdateEntrySet {utDate})) = utDate
binDate (ToRead ReadEntry {reDate}) = reDate
binDate (ToInsert (Tx {txDate})) = txDate

type EntryBals = M.Map (AccountRId, CurrencyRId) Rational

data UpdateEntryType a
  = UEReadOnly (UpdateEntry () Rational)
  | UEBlank (UpdateEntry EntryRId Rational)
  | UEPaired (UpdateEntry EntryRId Rational, UpdateEntry EntryRId a)

rebalanceEntrySet :: UpdateEntrySet -> State EntryBals [UpdateEntry EntryRId Rational]
rebalanceEntrySet
  UpdateEntrySet
    { utFrom0
    , utTo0
    , utPairs
    , utFromUnk
    , utToUnk
    , utFromRO
    , utToRO
    , utCurrency
    , utTotalValue
    } =
    do
      let fs =
            L.sortOn index $
              (UEReadOnly <$> utFromRO)
                ++ (UEBlank <$> utFromUnk)
                ++ (UEPaired <$> utPairs)
      fs' <- mapM goFrom fs
      let f0 = utFrom0 {ueValue = utTotalValue - (sum $ fmap value fs')}
      let (fs'', tpairs) = partitionEithers $ concatMap flatten fs'
      let ts = (Right <$> tpairs) ++ (Right <$> utToUnk) ++ (Left <$> utToRO)
      (tsRO, tsUnk) <- partitionEithers <$> mapM goTo ts
      let t0 = utTo0 {ueValue = utTotalValue - (sum $ (fmap ueValue tsRO) ++ (fmap ueValue tsUnk))}
      return $ f0 : fs'' ++ t0 : tsUnk
    where
      project f _ _ (UEReadOnly e) = f e
      project _ f _ (UEBlank e) = f e
      project _ _ f (UEPaired p) = f p
      index = project ueIndex ueIndex (ueIndex . fst)
      value = project ueValue ueValue (ueValue . fst)
      flatten = project (const []) ((: []) . Right) (\(a, b) -> [Right a, Left b])
      -- TODO the following is wetter than the average groupie
      goFrom (UEReadOnly e) = do
        modify $ mapAdd_ (ueAcnt e, utCurrency) (ueValue e)
        return $ UEReadOnly e
      goFrom (UEBlank e) = do
        let key = (ueAcnt e, utCurrency)
        curBal <- gets (M.findWithDefault 0 key)
        let newVal = ueValue e - curBal
        modify $ mapAdd_ key newVal
        return $ UEBlank $ e {ueValue = newVal}
      goFrom (UEPaired (e0, e1)) = do
        let key = (ueAcnt e0, utCurrency)
        curBal <- gets (M.findWithDefault 0 key)
        let newVal = ueValue e0 - curBal
        modify $ mapAdd_ key newVal
        return $ UEPaired $ (e0 {ueValue = newVal}, e1 {ueValue = -newVal})
      goTo (Left e) = do
        modify $ mapAdd_ (ueAcnt e, utCurrency) (ueValue e)
        return $ Left e
      goTo (Right e) = do
        let key = (ueAcnt e, utCurrency)
        curBal <- gets (M.findWithDefault 0 key)
        let newVal = ueValue e - curBal
        modify $ mapAdd_ key newVal
        return $ Right $ e {ueValue = newVal}

balanceEntrySet
  :: (MonadInsertError m, MonadFinance m)
  => Day
  -> DeferredEntrySet
  -> StateT EntryBals m [KeyEntry]
balanceEntrySet
  day
  EntrySet
    { esFrom = HalfEntrySet {hesPrimary = f0, hesOther = fs}
    , esTo = HalfEntrySet {hesPrimary = t0, hesOther = ts}
    , esCurrency
    , esTotalValue
    } =
    do
      -- get currency first and quit immediately on exception since everything
      -- downstream depends on this
      (curID, precision) <- lookupCurrency esCurrency

      -- resolve accounts and balance debit entries since we need an array
      -- of debit entries for linked credit entries later
      let balFromEntry = balanceEntry (balanceDeferred curID) curID
      fs' <- doEntries balFromEntry curID esTotalValue f0 fs (NE.iterate (-1) (-1))
      let fv = V.fromList $ fmap (eValue . feEntry) fs'

      -- finally resolve credit entries
      let balToEntry = balanceEntry (balanceLinked fv curID precision) curID
      ts' <- doEntries balToEntry curID (-esTotalValue) t0 ts (NE.iterate (+ 1) 0)
      return $ fs' ++ ts'

doEntries
  :: (MonadInsertError m, MonadFinance m)
  => (Int -> Entry AcntID v TagID -> State EntryBals (FullEntry AccountRId CurrencyRId TagID))
  -> CurrencyRId
  -> Rational
  -> Entry AcntID () TagID
  -> [Entry AcntID v TagID]
  -> NonEmpty Int
  -> StateT EntryBals m [FullEntry AccountRId CurrencyRId TagID]
doEntries f curID tot e es (i0 :| iN) = do
  es' <- liftInnerS $ mapM (uncurry f) $ zip iN es
  let val0 = tot - entrySum es'
  e' <- balanceEntry (\_ _ -> return (val0, Nothing)) curID i0 e
  return $ e' : es'
  where
    entrySum = sum . fmap (eValue . feEntry)

liftInnerS :: Monad m => StateT e Identity a -> StateT e m a
liftInnerS = mapStateT (return . runIdentity)

balanceLinked
  :: Vector Rational
  -> CurrencyRId
  -> Natural
  -> AccountRId
  -> LinkDeferred Rational
  -> StateT EntryBals Identity (Rational, Maybe DBDeferred)
balanceLinked from curID precision acntID lg = case lg of
  (LinkIndex g@LinkedNumGetter {lngIndex, lngScale}) -> do
    let res = fmap (go lngScale) $ from V.!? fromIntegral lngIndex
    case res of
      Just v -> return $ (v, Just $ EntryLinked lngIndex $ toRational lngScale)
      Nothing -> throwError undefined
  (LinkDeferred d) -> balanceDeferred curID acntID d
  where
    go s = roundPrecision precision . (* s) . fromRational

balanceDeferred
  :: CurrencyRId
  -> AccountRId
  -> Deferred Rational
  -> State EntryBals (Rational, Maybe DBDeferred)
balanceDeferred curID acntID (Deferred toBal v) = do
  newval <- findBalance acntID curID toBal v
  return $ (newval, if toBal then Just (EntryBalance v) else Nothing)

balanceEntry
  :: (MonadInsertError m, MonadFinance m)
  => (AccountRId -> v -> StateT EntryBals m (Rational, Maybe DBDeferred))
  -> CurrencyRId
  -> Int
  -> Entry AcntID v TagID
  -> StateT EntryBals m (FullEntry AccountRId CurrencyRId TagID)
balanceEntry f curID index e@Entry {eValue, eAcnt} = do
  (acntID, sign, _) <- lookupAccount eAcnt
  let s = fromIntegral $ sign2Int sign
  (newVal, deferred) <- f acntID eValue
  modify (mapAdd_ (acntID, curID) newVal)
  return $
    FullEntry
      { feEntry = e {eValue = s * newVal, eAcnt = acntID}
      , feCurrency = curID
      , feDeferred = deferred
      , feIndex = index
      }
  where
    key = (eAcnt, curID)

findBalance :: AccountRId -> CurrencyRId -> Bool -> Rational -> State EntryBals Rational
findBalance acnt cur toBal v = do
  curBal <- gets (M.findWithDefault 0 (acnt, cur))
  return $ if toBal then v - curBal else v

-- -- 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