pwncash/lib/Internal/Insert.hs

module Internal.Insert
  ( insertStatements
  , insertBudget
  )
where

import Data.Hashable
import Database.Persist.Class
import Database.Persist.Sql hiding (Single, Statement)
import GHC.Utils.Misc hiding (split)
import Internal.Statement
import Internal.Types hiding (sign)
import Internal.Utils
import RIO hiding (to)
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

--------------------------------------------------------------------------------
-- intervals

expandDatePat :: Bounds -> DatePat -> EitherErrs [Day]
expandDatePat b (Cron cp) = expandCronPat b cp
expandDatePat i (Mod mp) = Right $ expandModPat mp i

expandModPat :: ModPat -> Bounds -> [Day]
expandModPat ModPat {mpStart = s, mpBy = b, mpUnit = u, mpRepeats = r} bs =
  takeWhile (<= upper) $
    (`addFun` start) . (* b')
      <$> maybe id (take . fromIntegral) r [0 ..]
  where
    (lower, upper) = expandBounds bs
    start = maybe lower fromGregorian' s
    b' = fromIntegral b
    addFun = case u of
      Day -> addDays
      Week -> addDays . (* 7)
      Month -> addGregorianMonthsClip
      Year -> addGregorianYearsClip

expandCronPat :: Bounds -> CronPat -> EitherErrs [Day]
expandCronPat b CronPat {cronYear, cronMonth, cronDay, cronWeekly} =
  concatEither3 yRes mRes dRes $ \ys ms ds ->
    filter validWeekday $
      mapMaybe (uncurry3 toDay) $
        takeWhile (\((y, _), m, d) -> (y, m, d) <= (yb1, mb1, db1)) $
          dropWhile (\((y, _), m, d) -> (y, m, d) < (yb0, mb0, db0)) $
            [(y, m, d) | y <- (\y -> (y, isLeapYear y)) <$> ys, m <- ms, d <- ds]
  where
    yRes = case cronYear of
      Nothing -> return [yb0 .. yb1]
      Just pat -> do
        ys <- expandMDYPat (fromIntegral yb0) (fromIntegral yb1) pat
        return $ dropWhile (< yb0) $ fromIntegral <$> ys
    mRes = expandMD 12 cronMonth
    dRes = expandMD 31 cronDay
    (s, e) = expandBounds b
    (yb0, mb0, db0) = toGregorian s
    (yb1, mb1, db1) = toGregorian $ addDays (-1) e
    expandMD lim =
      fmap (fromIntegral <$>)
        . maybe (return [1 .. lim]) (expandMDYPat 1 lim)
    expandW (OnDay x) = [fromEnum x]
    expandW (OnDays xs) = fromEnum <$> xs
    ws = maybe [] expandW cronWeekly
    validWeekday = if null ws then const True else \day -> dayToWeekday day `elem` ws
    toDay (y, leap) m d
      | m == 2 && (not leap && d > 28 || leap && d > 29) = Nothing
      | m `elem` [4, 6, 9, 11] && d > 30 = Nothing
      | otherwise = Just $ fromGregorian y m d

expandMDYPat :: Natural -> Natural -> MDYPat -> EitherErr [Natural]
expandMDYPat lower upper (Single x) = Right [x | lower <= x && x <= upper]
expandMDYPat lower upper (Multi xs) = Right $ dropWhile (<= lower) $ takeWhile (<= upper) xs
expandMDYPat lower upper (After x) = Right [max lower x .. upper]
expandMDYPat lower upper (Before x) = Right [lower .. min upper x]
expandMDYPat lower upper (Between x y) = Right [max lower x .. min upper y]
expandMDYPat lower upper (Repeat RepeatPat {rpStart = s, rpBy = b, rpRepeats = r})
  | b < 1 = Left $ PatternError s b r ZeroLength
  | otherwise = do
      k <- limit r
      return $ dropWhile (<= lower) $ takeWhile (<= k) [s + i * b | i <- [0 ..]]
  where
    limit Nothing = Right upper
    limit (Just n)
      -- this guard not only produces the error for the user but also protects
      -- from an underflow below it
      | n < 1 = Left $ PatternError s b r ZeroRepeats
      | otherwise = Right $ min (s + b * (n - 1)) upper

dayToWeekday :: Day -> Int
dayToWeekday (ModifiedJulianDay d) = mod (fromIntegral d + 2) 7

withDates
  :: MonadFinance m
  => DatePat
  -> (Day -> SqlPersistT m (EitherErrs a))
  -> SqlPersistT m (EitherErrs [a])
withDates dp f = do
  bounds <- lift $ askDBState kmBudgetInterval
  case expandDatePat bounds dp of
    Left es -> return $ Left es
    Right days -> concatEithersL <$> mapM f days

--------------------------------------------------------------------------------
-- budget

-- each budget (designated at the top level by a 'name') is processed in the
-- following steps
-- 1. expand all transactions given the desired date range and date patterns for
--    each directive in the budget
-- 2. sort all transactions by date
-- 3. propagate all balances forward, and while doing so assign values to each
--    transaction (some of which depend on the 'current' balance of the
--    target account)
-- 4. assign shadow transactions (TODO)
-- 5. insert all transactions

insertBudget :: MonadFinance m => Budget -> SqlPersistT m [InsertError]
insertBudget Budget {budgetLabel = name, income = is, transfers = es, shadowTransfers = ss} = do
  res1 <- mapM (insertIncome name) is
  res2 <- expandTransfers name es
  unlessLefts (concatEithers2 (fmap concat $ concatEithersL res1) res2 (++)) $
    \txs -> do
      unlessLefts (addShadowTransfers ss txs) $ \shadow -> do
        let bals = balanceTransfers $ txs ++ shadow
        concat <$> mapM insertBudgetTx bals

-- TODO this is going to be O(n*m), which might be a problem?
addShadowTransfers :: [ShadowTransfer] -> [BudgetTxType] -> EitherErrs [BudgetTxType]
addShadowTransfers ms txs =
  fmap catMaybes $
    concatEitherL $
      fmap (uncurry fromShadow) $
        [(t, m) | t <- txs, m <- ms]

fromShadow :: BudgetTxType -> ShadowTransfer -> EitherErr (Maybe BudgetTxType)
fromShadow tx t@ShadowTransfer {stFrom, stTo, stDesc, stRatio} = do
  res <- shadowMatches (stMatch t) tx
  return $
    if not res
      then Nothing
      else
        Just $
          BudgetTxType
            { bttTx =
                -- TODO does this actually share the same metadata as the "parent" tx?
                BudgetTx
                  { btMeta = btMeta $ bttTx tx
                  , btWhen = btWhen $ bttTx tx
                  , btFrom = stFrom
                  , btTo = stTo
                  , btValue = dec2Rat stRatio * (btValue $ bttTx tx)
                  , btDesc = stDesc
                  }
            , bttType = FixedAmt
            }

shadowMatches :: ShadowMatch -> BudgetTxType -> EitherErr Bool
shadowMatches ShadowMatch {smFrom, smTo, smDate, smVal} tx = do
  -- TODO what does the amount do for each of the different types?
  valRes <- valMatches smVal (btValue tx_)
  return $
    memberMaybe (btFrom tx_) smFrom
      && memberMaybe (btTo tx_) smTo
      && maybe True (`dateMatches` (btWhen tx_)) smDate
      && valRes
  where
    tx_ = bttTx tx
    memberMaybe x AcntSet {asList, asInclude} =
      (if asInclude then id else not) $ x `elem` asList

balanceTransfers :: [BudgetTxType] -> [BudgetTx]
balanceTransfers ts = snd $ L.mapAccumR go initBals $ L.sortOn (btWhen . bttTx) ts
  where
    initBals =
      M.fromList $
        fmap (,0) $
          L.nub $
            (fmap (btTo . bttTx) ts ++ fmap (btTo . bttTx) ts)
    updateBal x = M.update (Just . (+ x))
    lookupBal = M.findWithDefault (error "this should not happen")
    go bals btt =
      let tx = bttTx btt
          from = btFrom tx
          to = btTo tx
          bal = lookupBal to bals
          x = amtToMove bal (bttType btt) (btValue tx)
       in (updateBal x to $ updateBal (-x) from bals, tx {btValue = x})
    -- TODO might need to query signs to make this intuitive; as it is this will
    -- probably work, but for credit accounts I might need to supply a negative
    -- target value
    amtToMove _ FixedAmt x = x
    amtToMove bal Percent x = -(x / 100 * bal)
    amtToMove bal Target x = x - bal

data BudgetMeta = BudgetMeta
  { bmCommit :: !(Key CommitR)
  , bmCur :: !BudgetCurrency
  , bmName :: !T.Text
  }

data BudgetTx = BudgetTx
  { btMeta :: !BudgetMeta
  , btWhen :: !Day
  , btFrom :: !AcntID
  , btTo :: !AcntID
  , btValue :: !Rational
  , btDesc :: !T.Text
  }

data BudgetTxType = BudgetTxType
  { bttType :: !AmountType
  , bttTx :: !BudgetTx
  }

insertIncome :: MonadFinance m => T.Text -> Income -> SqlPersistT m (EitherErrs [BudgetTxType])
insertIncome
  name
  i@Income {incWhen, incCurrency, incFrom, incPretax, incPosttax, incTaxes, incToBal} =
    whenHash CTIncome i (Right []) $ \c -> do
      let meta = BudgetMeta c (NoX incCurrency) name
      let balRes = balanceIncome i
      fromRes <- lift $ checkAcntType IncomeT incFrom
      case concatEither2 balRes fromRes (,) of
        Left es -> return $ Left es
        Right (balance, from) ->
          fmap (fmap (concat . concat)) $
            withDates incWhen $ \day -> do
              let fromAllos = fmap concat . mapM (lift . fromAllo day meta from)
              pre <- fromAllos incPretax
              tax <-
                concatEitherL
                  <$> mapM (lift . fromTax day meta from) incTaxes
              post <- fromAllos incPosttax
              let bal =
                    BudgetTxType
                      { bttTx =
                          BudgetTx
                            { btMeta = meta
                            , btWhen = day
                            , btFrom = from
                            , btTo = incToBal
                            , btValue = balance
                            , btDesc = "balance after deductions"
                            }
                      , bttType = FixedAmt
                      }
              return $ concatEithersL [Right [bal], tax, Right pre, Right post]

fromAllo
  :: MonadFinance m
  => Day
  -> BudgetMeta
  -> AcntID
  -> Allocation
  -> m [BudgetTxType]
fromAllo day meta from Allocation {alloPath, alloAmts} = do
  -- TODO this is going to be repeated a zillion times (might matter)
  -- res <- expandTarget alloPath
  return $ fmap toBT alloAmts
  where
    toBT (Amount desc v) =
      BudgetTxType
        { bttTx =
            BudgetTx
              { btFrom = from
              , btWhen = day
              , btTo = alloPath
              , btValue = dec2Rat v
              , btDesc = desc
              , btMeta = meta
              }
        , bttType = FixedAmt
        }

fromTax
  :: MonadFinance m
  => Day
  -> BudgetMeta
  -> AcntID
  -> Tax
  -> m (EitherErr BudgetTxType)
fromTax day meta from Tax {taxAcnt = to, taxValue = v} = do
  res <- checkAcntType ExpenseT to
  return $ fmap go res
  where
    go to_ =
      BudgetTxType
        { bttTx =
            BudgetTx
              { btFrom = from
              , btWhen = day
              , btTo = to_
              , btValue = dec2Rat v
              , btDesc = ""
              , btMeta = meta
              }
        , bttType = FixedAmt
        }

balanceIncome :: Income -> EitherErr Rational
balanceIncome
  Income
    { incGross = g
    , incWhen = dp
    , incPretax = pre
    , incTaxes = tax
    , incPosttax = post
    }
    | bal < 0 = Left $ IncomeError dp
    | otherwise = Right bal
    where
      bal = dec2Rat g - sum (sumAllocation <$> pre ++ post) - sumTaxes tax

sumAllocation :: Allocation -> Rational
sumAllocation = sum . fmap (dec2Rat . amtValue) . alloAmts

sumTaxes :: [Tax] -> Rational
sumTaxes = sum . fmap (dec2Rat . taxValue)

expandTransfers
  :: MonadFinance m
  => T.Text
  -> [Transfer]
  -> SqlPersistT m (EitherErrs [BudgetTxType])
expandTransfers name ts = do
  txs <- mapM (expandTransfer name) ts
  return $ L.sortOn (btWhen . bttTx) . concat <$> concatEithersL txs

expandTransfer :: MonadFinance m => T.Text -> Transfer -> SqlPersistT m (EitherErrs [BudgetTxType])
expandTransfer name t@Transfer {transAmounts, transTo, transCurrency, transFrom} =
  whenHash CTExpense t (Right []) $ \key ->
    fmap (fmap concat . concatEithersL) $
      forM transAmounts $ \(TimeAmount (Amount desc v) atype pat) -> do
        withDates pat $ \day ->
          let meta =
                BudgetMeta
                  { bmCur = transCurrency
                  , bmCommit = key
                  , bmName = name
                  }
              tx =
                BudgetTxType
                  { bttTx =
                      BudgetTx
                        { btMeta = meta
                        , btWhen = day
                        , btFrom = transFrom
                        , btTo = transTo
                        , btValue = dec2Rat v
                        , btDesc = desc
                        }
                  , bttType = atype
                  }
           in return $ Right tx

insertBudgetTx :: MonadFinance m => BudgetTx -> SqlPersistT m [InsertError]
insertBudgetTx BudgetTx {btFrom, btTo, btMeta, btValue, btDesc, btWhen} = do
  res <- lift $ splitPair btFrom btTo (bmCur btMeta) btValue
  unlessLefts_ res $ \((sFrom, sTo), exchange) -> do
    insertPair sFrom sTo
    forM_ exchange $ \(xFrom, xTo) -> insertPair xFrom xTo
  where
    insertPair from to = do
      k <- insert $ TransactionR (bmCommit btMeta) btWhen btDesc
      insertBudgetLabel k from
      insertBudgetLabel k to
    insertBudgetLabel k split = do
      sk <- insertSplit k split
      insert_ $ BudgetLabelR sk $ bmName btMeta

type SplitPair = (KeySplit, KeySplit)

splitPair
  :: MonadFinance m
  => AcntID
  -> AcntID
  -> BudgetCurrency
  -> Rational
  -> m (EitherErrs (SplitPair, Maybe SplitPair))
splitPair from to cur val = case cur of
  NoX curid -> fmap (fmap (,Nothing)) $ pair curid from to val
  X (Exchange {xFromCur, xToCur, xAcnt, xRate}) -> do
    res1 <- pair xFromCur from xAcnt val
    res2 <- pair xToCur xAcnt to (val * dec2Rat xRate)
    return $ concatEithers2 res1 res2 $ \a b -> (a, Just b)
  where
    pair curid from_ to_ v = do
      s1 <- split curid from_ (-v)
      s2 <- split curid to_ v
      return $ concatEithers2 s1 s2 (,)
    split c a v =
      resolveSplit $
        Split
          { sAcnt = a
          , sValue = v
          , sComment = ""
          , sCurrency = c
          }

checkAcntType
  :: MonadFinance m
  => AcntType
  -> AcntID
  -> m (EitherErr AcntID)
checkAcntType t = checkAcntTypes (t :| [])

checkAcntTypes
  :: MonadFinance m
  => NE.NonEmpty AcntType
  -> AcntID
  -> m (EitherErr AcntID)
checkAcntTypes ts i = (go =<<) <$> lookupAccountType i
  where
    go t
      | t `L.elem` ts = Right i
      | otherwise = Left $ AccountError i ts

--------------------------------------------------------------------------------
-- statements

insertStatements :: MonadFinance m => Config -> SqlPersistT m [InsertError]
insertStatements conf = concat <$> mapM insertStatement (statements conf)

insertStatement :: MonadFinance m => Statement -> SqlPersistT m [InsertError]
insertStatement (StmtManual m) = insertManual m
insertStatement (StmtImport i) = insertImport i

insertManual :: MonadFinance m => Manual -> SqlPersistT m [InsertError]
insertManual
  m@Manual
    { manualDate = dp
    , manualFrom = from
    , manualTo = to
    , manualValue = v
    , manualCurrency = u
    , manualDesc = e
    } = do
    whenHash CTManual m [] $ \c -> do
      bounds <- lift $ askDBState kmStatementInterval
      -- let days = expandDatePat bounds dp
      let dayRes = expandDatePat bounds dp
      unlessLefts dayRes $ \days -> do
        txRes <- mapM (lift . tx) days
        unlessLefts_ (concatEithersL txRes) $ mapM_ (insertTx c)
    where
      tx day = txPair day from to u (dec2Rat v) e

insertImport :: MonadFinance m => Import -> SqlPersistT m [InsertError]
insertImport i = whenHash CTImport i [] $ \c -> do
  -- TODO this isn't efficient, the whole file will be read and maybe no
  -- transactions will be desired
  recoverIO (lift $ readImport i) $ \r -> unlessLefts r $ \bs -> do
    bounds <- expandBounds <$> lift (askDBState kmStatementInterval)
    res <- mapM (lift . resolveTx) $ filter (inBounds bounds . txDate) bs
    unlessLefts_ (concatEithersL res) $ mapM_ (insertTx c)
  where
    recoverIO x rest = do
      res <- tryIO x
      case res of
        Right r -> rest r
        -- If file is not found (or something else happens) then collect the
        -- error try the remaining imports
        Left e -> return [InsertIOError $ showT e]

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

txPair
  :: MonadFinance m
  => Day
  -> AcntID
  -> AcntID
  -> CurID
  -> Rational
  -> T.Text
  -> m (EitherErrs KeyTx)
txPair day from to cur val desc = resolveTx tx
  where
    split a v = Split {sAcnt = a, sValue = v, sComment = "", sCurrency = cur}
    tx =
      Tx
        { txDescr = desc
        , txDate = day
        , txTags = []
        , txSplits = [split from (-val), split to val]
        }

resolveTx :: MonadFinance m => BalTx -> m (EitherErrs KeyTx)
resolveTx t@Tx {txSplits = ss} = do
  res <- concatEithersL <$> mapM resolveSplit ss
  return $ fmap (\kss -> t {txSplits = kss}) res

resolveSplit :: MonadFinance m => BalSplit -> m (EitherErrs KeySplit)
resolveSplit s@Split {sAcnt = p, sCurrency = c, sValue = v} = do
  aid <- lookupAccountKey p
  cid <- lookupCurrency c
  sign <- lookupAccountSign p
  -- TODO correct sign here?
  -- TODO lenses would be nice here
  return $ concatEither3 aid cid sign $ \aid_ cid_ sign_ ->
    s
      { sAcnt = aid_
      , sCurrency = cid_
      , sValue = v * fromIntegral (sign2Int sign_)
      }

insertTx :: MonadUnliftIO m => Key CommitR -> KeyTx -> SqlPersistT m ()
insertTx c Tx {txDate = d, txDescr = e, txSplits = ss} = do
  k <- insert $ TransactionR c d e
  mapM_ (insertSplit k) ss

insertSplit :: MonadUnliftIO m => Key TransactionR -> KeySplit -> SqlPersistT m (Key SplitR)
insertSplit t Split {sAcnt = aid, sCurrency = cid, sValue = v, sComment = c} = do
  insert $ SplitR t cid aid c v

lookupAccount :: MonadFinance m => AcntID -> m (EitherErr (Key AccountR, AcntSign, AcntType))
lookupAccount p = lookupErr (DBKey AcntField) p <$> (askDBState kmAccount)

lookupAccountKey :: MonadFinance m => AcntID -> m (EitherErr (Key AccountR))
lookupAccountKey = (fmap (fmap fstOf3)) . lookupAccount

lookupAccountSign :: MonadFinance m => AcntID -> m (EitherErr AcntSign)
lookupAccountSign = fmap (fmap sndOf3) . lookupAccount

lookupAccountType :: MonadFinance m => AcntID -> m (EitherErr AcntType)
lookupAccountType = fmap (fmap thdOf3) . lookupAccount

lookupCurrency :: MonadFinance m => T.Text -> m (EitherErr (Key CurrencyR))
lookupCurrency c = lookupErr (DBKey CurField) c <$> (askDBState kmCurrency)

-- TODO this hashes twice (not that it really matters)
whenHash
  :: (Hashable a, MonadFinance m)
  => ConfigType
  -> a
  -> b
  -> (Key CommitR -> SqlPersistT m b)
  -> SqlPersistT m b
whenHash t o def f = do
  let h = hash o
  hs <- lift $ askDBState kmNewCommits
  if h `elem` hs then f =<< insert (CommitR h t) else return def