pwncash/lib/Internal/Insert.hs

{-# LANGUAGE GADTs #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE NoImplicitPrelude #-}

module Internal.Insert
  ( insertStatements
  , insertBudget
  )
where

import Data.Hashable
import Database.Persist.Class
import Database.Persist.Sql hiding (Single, Statement)
import Internal.Database.Model
import Internal.Statement
import Internal.Types hiding (sign)
import Internal.Utils
import RIO hiding (to)
import qualified RIO.Text as T
import RIO.Time

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

expandDatePat :: Bounds -> DatePat -> [Day]
expandDatePat (a, b) (Cron cp) = filter (cronPatternMatches cp) [a .. b]
expandDatePat i (Mod mp) = expandModPat mp i

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

-- TODO this can be optimized to prevent filtering a bunch of dates for
-- one/a few cron patterns
cronPatternMatches :: CronPat -> Day -> Bool
cronPatternMatches
  CronPat
    { cronWeekly = w
    , cronYear = y
    , cronMonth = m
    , cronDay = d
    }
  x =
    yMaybe (y' - 2000) y && mdMaybe m' m && mdMaybe d' d && wdMaybe (dayOfWeek_ x) w
    where
      testMaybe = maybe True
      yMaybe z = testMaybe (mdyPatternMatches testYear (fromIntegral z))
      mdMaybe z = testMaybe (mdyPatternMatches (const Nothing) (fromIntegral z))
      wdMaybe z = testMaybe (`weekdayPatternMatches` z)
      (y', m', d') = toGregorian x
      testYear z = if z > 99 then Just "year must be 2 digits" else Nothing

dayOfWeek_ :: Day -> Weekday
dayOfWeek_ d = case dayOfWeek d of
  Monday -> Mon
  Tuesday -> Tue
  Wednesday -> Wed
  Thursday -> Thu
  Friday -> Fri
  Saturday -> Sat
  Sunday -> Sun

weekdayPatternMatches :: WeekdayPat -> Weekday -> Bool
weekdayPatternMatches (OnDay x) = (== x)
weekdayPatternMatches (OnDays xs) = (`elem` xs)

mdyPatternMatches :: (Natural -> Maybe String) -> Natural -> MDYPat -> Bool
mdyPatternMatches check x p = case p of
  Single y -> errMaybe (check y) $ x == y
  Multi xs -> errMaybe (msum $ check <$> xs) $ x `elem` xs
  Repeat (RepeatPat {rpStart = s, rpBy = b, rpRepeats = r}) ->
    errMaybe (check s) $
      s >= x && mod x b == 0 && maybe True (\y -> x <= s + b * y) r
  where
    errMaybe test rest = maybe rest err test
    err msg = error $ show p ++ ": " ++ msg

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

insertBudget :: MonadUnliftIO m => Budget -> MappingT m [InsertError]
insertBudget Budget {income = is, expenses = es} = do
  es1 <- mapM insertIncome is
  es2 <- mapM insertExpense es
  return $ concat $ es1 ++ es2

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

insertIncome :: MonadUnliftIO m => Income -> MappingT m [InsertError]
insertIncome
  i@Income
    { incCurrency = cur
    , incWhen = dp
    , incAccount = from
    , incTaxes = taxes
    } =
    whenHash CTIncome i [] $ \c ->
      unlessLeft (balanceIncome i) $ \balanced -> do
        bounds <- (liftIO . resolveBounds) =<< asks kmBudgetInterval
        fmap concat $ forM (expandDatePat bounds dp) $ \day -> do
          -- TODO why are these separate?
          nontaxRes <- alloTxs concat (allocationToTx from day) balanced
          taxRes <- alloTxs (fmap (,Fixed)) (taxToTx from day cur) taxes
          unlessLefts_ (concatEithers2 nontaxRes taxRes (++)) $ \txs ->
            lift $ mapM_ (\(t, b) -> insertTxBucket (Just b) c t) txs
    where
      alloTxs squish toTx = fmap (fmap squish . concatEithersL) . mapM toTx

balanceIncome :: Income -> EitherErr [BalAllocation]
balanceIncome
  Income
    { incGross = g
    , incWhen = dp
    , incPretax = pre
    , incTaxes = tax
    , incPosttax = post
    } = (preRat ++) <$> balancePostTax dp bal postRat
    where
      preRat = mapAlloAmts dec2Rat <$> pre
      postRat = mapAlloAmts (fmap dec2Rat) <$> post
      bal = dec2Rat g - (sumAllocations preRat + sumTaxes tax)

mapAlloAmts :: (a -> b) -> Allocation a -> Allocation b
mapAlloAmts f a@Allocation {alloAmts = as} = a {alloAmts = fmap f <$> as}

sumAllocations :: [BalAllocation] -> Rational
sumAllocations = sum . concatMap (fmap amtValue . alloAmts)

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

balancePostTax :: DatePat -> Rational -> [RawAllocation] -> EitherErr [BalAllocation]
balancePostTax dp bal as
  | null as = err NoAllocations
  | otherwise = case partitionEithers $ fmap hasVal as of
      ([([empty], nonmissing)], bs) ->
        let s = bal - sumAllocations (nonmissing : bs)
         in if s < 0
              then err ExceededTotal
              else Right $ mapAmts (empty {amtValue = s} :) nonmissing : bs
      ([], _) -> err MissingBlank
      _ -> err TooManyBlanks
  where
    hasVal a@Allocation {alloAmts = xs} =
      case partitionEithers $ fmap maybeAmt xs of
        ([], bs) -> Right a {alloAmts = bs}
        (unbal, bs) -> Left (unbal, a {alloAmts = bs})
    maybeAmt a@Amount {amtValue = Just v} = Right a {amtValue = v}
    maybeAmt a = Left a
    err t = Left $ AllocationError t dp

-- TODO lens reinvention
mapAmts :: ([Amount a] -> [Amount b]) -> Allocation a -> Allocation b
mapAmts f a@Allocation {alloAmts = xs} = a {alloAmts = f xs}

allocationToTx
  :: MonadUnliftIO m
  => AcntID
  -> Day
  -> BalAllocation
  -> MappingT m (EitherErrs [(KeyTx, Bucket)])
allocationToTx
  from
  day
  Allocation
    { alloPath = to
    , alloBucket = b
    , alloCurrency = cur
    , alloAmts = as
    } =
    second (fmap (,b)) . concatEithersL <$> mapM (transferToTx day from to cur) as

taxToTx
  :: MonadUnliftIO m
  => AcntID
  -> Day
  -> T.Text
  -> Tax
  -> MappingT m (EitherErrs KeyTx)
taxToTx from day cur Tax {taxAcnt = to, taxValue = v} =
  txPair day from to cur (dec2Rat v) ""

transferToTx
  :: MonadUnliftIO m
  => Day
  -> AcntID
  -> AcntID
  -> T.Text
  -> BalAmount
  -> MappingT m (EitherErrs KeyTx)
transferToTx day from to cur Amount {amtValue = v, amtDesc = d} =
  txPair day from to cur v d

insertExpense :: MonadUnliftIO m => Expense -> MappingT m [InsertError]
insertExpense
  e@Expense
    { expFrom = from
    , expTo = to
    , expCurrency = cur
    , expBucket = buc
    , expAmounts = as
    } = do
    whenHash CTExpense e [] $ \key -> concat <$> mapM (go key) as
    where
      go key amt = do
        res <- timeAmountToTx from to cur amt
        unlessLefts_ res $
          lift . mapM_ (insertTxBucket (Just buc) key)

timeAmountToTx
  :: MonadUnliftIO m
  => AcntID
  -> AcntID
  -> CurID
  -> TimeAmount
  -> MappingT m (EitherErrs [KeyTx])
timeAmountToTx
  from
  to
  cur
  TimeAmount
    { taWhen = dp
    , taAmt =
      Amount
        { amtValue = v
        , amtDesc = d
        }
    } = do
    bounds <- (liftIO . resolveBounds) =<< asks kmBudgetInterval
    concatEithersL <$> mapM tx (expandDatePat bounds dp)
    where
      tx day = txPair day from to cur (dec2Rat v) d

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

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

-- unless (null es) $ throwIO $ InsertException es

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

insertManual :: MonadUnliftIO m => Manual -> MappingT m [InsertError]
insertManual
  m@Manual
    { manualDate = dp
    , manualFrom = from
    , manualTo = to
    , manualValue = v
    , manualCurrency = u
    , manualDesc = e
    } = do
    whenHash CTManual m [] $ \c -> do
      bounds <- (liftIO . resolveBounds) =<< asks kmStatementInterval
      res <- mapM tx $ expandDatePat bounds dp
      unlessLefts_ (concatEithersL res) $ lift . mapM_ (insertTx c)
    where
      tx day = txPair day from to u (dec2Rat v) e

insertImport :: MonadUnliftIO m => Import -> MappingT 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 (readImport i) $ \r -> unlessLefts r $ \bs -> do
    bounds <- asks kmStatementInterval
    res <- mapM resolveTx $ filter (inMaybeBounds bounds . txDate) bs
    unlessLefts_ (concatEithersL res) $ lift . 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
  :: MonadUnliftIO m
  => Day
  -> AcntID
  -> AcntID
  -> T.Text
  -> Rational
  -> T.Text
  -> MappingT 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 :: MonadUnliftIO m => BalTx -> MappingT m (EitherErrs KeyTx)
resolveTx t@Tx {txSplits = ss} = do
  res <- concatEithersL <$> mapM resolveSplit ss
  return $ fmap (\kss -> t {txSplits = kss}) res

resolveSplit :: MonadUnliftIO m => BalSplit -> MappingT 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_)
      }

-- return $ case (aid, cid, sign) of
--   _ -> Nothing

insertTxBucket :: MonadUnliftIO m => Maybe Bucket -> Key CommitR -> KeyTx -> SqlPersistT m ()
insertTxBucket b c Tx {txDate = d, txDescr = e, txSplits = ss} = do
  k <- insert $ TransactionR c d e (fmap (T.pack . show) b)
  mapM_ (insertSplit k) ss

insertTx :: MonadUnliftIO m => Key CommitR -> KeyTx -> SqlPersistT m ()
insertTx = insertTxBucket Nothing

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

lookupAccount :: MonadUnliftIO m => AcntID -> MappingT m (EitherErr (Key AccountR, AcntSign))
lookupAccount p = lookupErr (DBKey AcntField) p <$> asks kmAccount

lookupAccountKey :: MonadUnliftIO m => AcntID -> MappingT m (EitherErr (Key AccountR))
lookupAccountKey = fmap (fmap fst) . lookupAccount

lookupAccountSign :: MonadUnliftIO m => AcntID -> MappingT m (EitherErr AcntSign)
lookupAccountSign = fmap (fmap snd) . lookupAccount

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