pwncash/lib/Internal/History.hs

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

import Control.Monad.Except
import Data.Csv
import Database.Persist.Monad
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 qualified RIO.Text as T
import RIO.Time
import qualified RIO.Vector as V

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, [RawTx])]
  -> 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

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)

readHistTransfer
  :: (MonadInsertError m, MonadFinance m)
  => HistTransfer
  -> m (Maybe (CommitR, [RawTx]))
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

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

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

-- TODO tags here?
txPair
  :: Day
  -> AcntID
  -> AcntID
  -> CurID
  -> Rational
  -> T.Text
  -> RawTx
txPair day from to cur val desc = tx
  where
    split a v =
      Entry
        { eAcnt = a
        , eValue = ConstD v
        , eComment = ""
        , eCurrency = cur
        , eTags = []
        }
    tx =
      Tx
        { txDescr = desc
        , txDate = day
        , txEntries = [split from (-val), split to val]
        }

resolveTx :: (MonadInsertError m, MonadFinance m) => BalTx -> m KeyTx
resolveTx t@Tx {txEntries = ss} =
  fmap (\kss -> t {txEntries = kss}) $
    combineErrors $
      fmap 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 [RawTx]
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 [RawTx]
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 RawTx)
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 RawTx)
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 ([RawTx], [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 ([RawTx], [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 ([RawTx], [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 ([RawTx], [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, RawTx)]
  -> m [(CommitR, KeyTx)]
balanceTxs ts = do
  bs <- mapM balanceTx $ snd $ L.mapAccumR balanceTxTargets M.empty ts'
  return $ zip cs bs
  where
    (cs, ts') = L.unzip $ L.sortOn (txDate . snd) ts

balanceTxTargets
  :: M.Map a Rational
  -> Tx (Entry a (Deferred Rational) c t)
  -> (M.Map a Rational, Tx (Entry a (Maybe Rational) c t))
balanceTxTargets = undefined

balanceEntryTargets
  :: M.Map a Rational
  -> Entry a (Deferred Rational) c t
  -> (M.Map a Rational, Entry a (Maybe Rational) c t)
balanceEntryTargets = undefined

balanceTx
  :: (MonadInsertError m, MonadFinance m)
  => Tx (Entry AcntID (Maybe Rational) CurID TagID)
  -> m KeyTx
balanceTx t@Tx {txEntries = ss} = do
  bs <- liftExcept $ balanceEntries ss
  resolveTx $ t {txEntries = bs}

balanceEntries :: [Entry AcntID (Maybe Rational) CurID TagID] -> InsertExcept [BalEntry]
balanceEntries ss =
  fmap concat
    <$> mapM (uncurry bal)
    $ groupByKey
    $ fmap (\s -> (eCurrency s, s)) ss
  where
    haeValue s@Entry {eValue = Just v} = Right s {eValue = v}
    haeValue s = Left s
    bal cur rss
      | length rss < 2 = throwError $ InsertException [BalanceError TooFewEntries cur rss]
      | otherwise = case partitionEithers $ fmap haeValue rss of
          ([noVal], val) -> return $ noVal {eValue = foldr (\s x -> x - eValue s) 0 val} : val
          ([], val) -> return val
          _ -> throwError $ InsertException [BalanceError NotOneBlank cur rss]

groupByKey :: Ord k => [(k, v)] -> [(k, [v])]
groupByKey = M.toList . M.fromListWith (++) . fmap (second (: []))