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.Database.Model
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} = do
  res1 <- mapM (insertIncome name) is
  res2 <- expandTransfers name es
  unlessLefts (concatEithers2 (fmap concat $ concatEithersL res1) res2 (++)) $
    \txs -> do
      let bals = balanceTransfers txs
      concat <$> mapM insertBudgetTx bals

balanceTransfers :: [BudgetTxType] -> [BudgetTx]
balanceTransfers ts = snd $ L.mapAccumR go initBals $ L.sortOn (btWhen . bttTx) ts
  where
    initBals =
      M.fromList $
        fmap (,0) $
          L.nub $
            (fmap (bsAcnt . btTo . bttTx) ts ++ fmap (bsAcnt . 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 = bsAcnt $ btFrom tx
          to = bsAcnt $ 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

-- TODO allow currency conversions here
data BudgetSplit b = BudgetSplit
  { bsAcnt :: !AcntID
  , bsBucket :: !(Maybe b)
  }

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

data BudgetTx = BudgetTx
  { btMeta :: !BudgetMeta
  , btWhen :: !Day
  , btFrom :: !(BudgetSplit IncomeBucket)
  , btTo :: !(BudgetSplit ExpenseBucket)
  , 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 incCurrency name
      let balRes = balanceIncome i
      fromRes <- lift $ checkAcntType IncomeT incFrom (\j -> BudgetSplit j . Just)
      toRes <- lift $ expandTarget incToBal
      case concatEither3 balRes fromRes toRes (,,) of
        Left es -> return $ Left es
        Right (balance, fromFun, to) ->
          fmap (fmap (concat . concat)) $
            withDates incWhen $ \day -> do
              let fromAllos b =
                    fmap (fmap concat . concatEitherL)
                      . mapM (lift . fromAllo day meta (fromFun b))
              pre <- fromAllos PreTax incPretax
              tax <-
                concatEitherL
                  <$> mapM (lift . fromTax day meta (fromFun IntraTax)) incTaxes
              post <- fromAllos PostTax incPosttax
              let bal =
                    BudgetTxType
                      { bttTx =
                          BudgetTx
                            { btMeta = meta
                            , btWhen = day
                            , btFrom = fromFun PostTax
                            , btTo = to
                            , btValue = balance
                            , btDesc = "balance after deductions"
                            }
                      , bttType = FixedAmt
                      }
              return $ concatEithersL [Right [bal], tax, pre, post]

fromAllo
  :: MonadFinance m
  => Day
  -> BudgetMeta
  -> BudgetSplit IncomeBucket
  -> Allocation
  -> m (EitherErr [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 $ (\to -> fmap (toBT to) alloAmts) <$> res
  where
    toBT to (Amount desc v) =
      BudgetTxType
        { bttTx =
            BudgetTx
              { btFrom = from
              , btWhen = day
              , btTo = to
              , btValue = dec2Rat v
              , btDesc = desc
              , btMeta = meta
              }
        , bttType = FixedAmt
        }

fromTax
  :: MonadFinance m
  => Day
  -> BudgetMeta
  -> BudgetSplit IncomeBucket
  -> Tax
  -> m (EitherErr BudgetTxType)
fromTax day meta from Tax {taxAcnt = to, taxValue = v} =
  -- TODO this is going to be repeated a zillion times (might matter)
  checkAcntType ExpenseT to $ \to_ ->
    BudgetTxType
      { bttTx =
          BudgetTx
            { btFrom = from
            , btWhen = day
            , btTo = BudgetSplit to_ (Just Fixed)
            , 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
        -- TODO this is going to be repeated a zillion times (might matter)
        res <- lift $ expandTarget transTo
        case res of
          Left e -> return $ Left [e]
          Right to -> withDates pat $ \day ->
            let meta =
                  BudgetMeta
                    { bmCur = transCurrency
                    , bmCommit = key
                    , bmName = name
                    }
                tx =
                  BudgetTxType
                    { bttTx =
                        BudgetTx
                          { btMeta = meta
                          , btWhen = day
                          , btFrom = BudgetSplit transFrom Nothing
                          , btTo = to
                          , 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 (bsAcnt btFrom) (bsAcnt btTo) (bmCur btMeta) btValue
  unlessLefts_ res $ \(sFrom, sTo) -> do
    k <- insert $ TransactionR (bmCommit btMeta) btWhen btDesc
    insertBudgetLabel name k IncomeBucketR sFrom btFrom
    insertBudgetLabel name k ExpenseBucketR sTo btTo
  where
    name = bmName btMeta

insertBudgetLabel
  :: (MonadUnliftIO m, PersistRecordBackend record SqlBackend)
  => T.Text
  -> Key TransactionR
  -> (Key BudgetLabelR -> a -> record)
  -> KeySplit
  -> BudgetSplit a
  -> SqlPersistT m ()
insertBudgetLabel name k bucketType split bs = do
  sk <- insertSplit k split
  bk <- insert $ BudgetLabelR sk name
  forM_ (bsBucket bs) $ insert_ . bucketType bk

splitPair
  :: MonadFinance m
  => AcntID
  -> AcntID
  -> CurID
  -> Rational
  -> m (EitherErrs (KeySplit, KeySplit))
splitPair from to cur val = do
  s1 <- split from (-val)
  s2 <- split to val
  return $ concatEithers2 s1 s2 (,)
  where
    split a v =
      resolveSplit $
        Split
          { sAcnt = a
          , sValue = v
          , sComment = ""
          , sCurrency = cur
          }

expandTarget
  :: MonadFinance m
  => TransferTarget
  -> m (EitherErr (BudgetSplit ExpenseBucket))
expandTarget t = case t of
  ExpenseTarget i b -> checkAcntType ExpenseT i $ (`BudgetSplit` (Just b))
  GenericTarget i ->
    checkAcntTypes (LiabilityT :| [AssetT, EquityT]) i $
      (`BudgetSplit` Nothing)

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

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

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