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 . txDate) 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
    { txDescr = desc
    , txDate = day
    , 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 -> m KeyTx
resolveTx t@Tx {txEntries = ss} =
  (\kss -> t {txEntries = kss}) <$> mapErrors resolveEntry ss

insertTx :: MonadSqlQuery m => CommitRId -> KeyTx -> m ()
insertTx c Tx {txDate = d, txDescr = e, txEntries = 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

balanceTxs
  :: (MonadInsertError m, MonadFinance m)
  => [(CommitR, DeferredTx)]
  -> m [(CommitR, KeyTx)]
balanceTxs ts = do
  keyts <- mapErrors resolveTx =<< evalStateT (mapM go ts') M.empty
  return $ zip cs keyts
  where
    (cs, ts') = L.unzip $ L.sortOn (txDate . snd) ts
    go t@Tx {txEntries, txDate} =
      (\es -> t {txEntries = concat es}) <$> mapM (balanceEntrySet txDate) txEntries

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
  :: (MonadInsertError m, MonadFinance m)
  => Day
  -> DeferredEntrySet
  -> StateT EntryBals m [BalEntry]
balanceEntrySet
  day
  EntrySet
    { esFrom = HalfEntrySet {hesPrimary = f0, hesOther = fs}
    , esTo = HalfEntrySet {hesPrimary = t0, hesOther = ts}
    , esCurrency
    , esTotalValue
    } =
    do
      let (lts, dts) = partitionEithers $ splitLinked <$> ts
      fs' <- doEntries fs esTotalValue f0
      let fv = V.fromList $ fmap eValue fs'
      lts' <- lift $ mapErrors (resolveLinked fv esCurrency day) lts
      ts' <- doEntries (dts ++ lts') (-esTotalValue) t0
      return $ toFull <$> fs' ++ ts'
    where
      doEntries es tot e0 = do
        es' <- liftInnerS $ mapM (balanceEntry esCurrency) es
        let val0 = tot - entrySum es'
        modify $ mapAdd_ (eAcnt e0, esCurrency) val0
        return $ e0 {eValue = val0} : es'
      toFull e = FullEntry {feEntry = e, feCurrency = esCurrency}
      splitLinked e@Entry {eValue} = case eValue of
        LinkIndex l -> Left e {eValue = l}
        LinkDeferred d -> Right e {eValue = d}
      liftInnerS = mapStateT (return . runIdentity)

resolveLinked
  :: (MonadInsertError m, MonadFinance m)
  => Vector Rational
  -> CurID
  -> Day
  -> Entry AcntID LinkedNumGetter TagID
  -> m (Entry AcntID (Deferred Rational) TagID)
resolveLinked from cur day e@Entry {eValue = LinkedNumGetter {lngIndex, lngScale}} = do
  curMap <- askDBState kmCurrency
  case from V.!? fromIntegral lngIndex of
    Nothing -> throwError $ InsertException [IndexError e day]
    Just v -> do
      v' <- liftExcept $ roundPrecisionCur cur curMap $ lngScale * fromRational v
      return $ e {eValue = Deferred False v'}

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

balanceEntry
  :: CurID
  -> Entry AcntID (Deferred Rational) TagID
  -> State EntryBals (Entry AcntID Rational TagID)
balanceEntry curID e@Entry {eValue = Deferred toBal v, eAcnt} = do
  curBal <- gets (M.findWithDefault 0 key)
  let newVal = if toBal then v - curBal else v
  modify (mapAdd_ key newVal)
  return $ e {eValue = newVal}
  where
    key = (eAcnt, curID)

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