pwncash/lib/Internal/Insert.hs

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

module Internal.Insert
  ( insertStatements
  , insertBudget
  ) where

import           Control.Monad
import           Control.Monad.IO.Class
import           Control.Monad.Trans.Class
import           Control.Monad.Trans.Reader

import           Data.Either
import           Data.Hashable
import qualified Data.Map                   as M
import           Data.Maybe
import qualified Data.Text                  as T
import           Data.Time

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           Numeric.Natural

lookupKey :: (Show v, Ord k, Show k, MonadIO m) => M.Map k v -> k -> m (Maybe v)
lookupKey m k = do
  let v = M.lookup k m
  when (isNothing v) $
    liftIO $ putStrLn $ "key does not exist: " ++ show k
  return v

lookupAccount :: MonadIO m => AcntID -> MappingT m (Maybe (Key AccountR, AcntSign))
lookupAccount p = do
  m <- asks kmAccount
  lookupKey m p

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

lookupAccountSign :: MonadIO m => AcntID -> MappingT m (Maybe AcntSign)
lookupAccountSign = fmap (fmap snd) . lookupAccount

lookupCurrency :: MonadIO m => T.Text -> MappingT m (Maybe (Key CurrencyR))
lookupCurrency c = do
  m <- asks kmCurrency
  lookupKey m c

--------------------------------------------------------------------------------
-- 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 :: MonadIO m => Budget -> MappingT m ()
insertBudget Budget { income = is, expenses = es } = do
  mapM_ insertIncome is
  mapM_ insertExpense es

-- TODO this hashes twice (not that it really matters)
whenHash :: Hashable a => MonadIO m => ConfigType -> a
  -> (Key CommitR -> MappingT m ()) -> MappingT m ()
whenHash t o f = do
  let h = hash o
  hs <- asks kmNewCommits
  when (h `elem` hs) $ do
    f =<< lift (insert $ CommitR h t)

insertIncome :: MonadIO m => Income -> MappingT m ()
insertIncome i@Income { incCurrency = cur
                      , incWhen = dp
                      , incAccount = from
                      , incTaxes = ts
                      } =
  whenHash CTIncome i $ \c -> do
  case balanceIncome i of
    Left m -> liftIO $ print m
    Right as -> do
      bounds <- (liftIO . resolveBounds) =<< asks kmBudgetInterval
      forM_ (expandDatePat bounds dp) $ \day -> do
        alloTx <- concat <$> mapM (allocationToTx from day) as
        taxTx <- fmap (, Fixed) <$> mapM (taxToTx from day cur) ts
        lift $ mapM_ (\(t, b) -> insertTxBucket (Just b) c t) $ alloTx ++ taxTx

balanceIncome :: Income -> Either T.Text [BalAllocation]
balanceIncome Income { incGross = g
                     , incPretax = pre
                     , incTaxes = tax
                     , incPosttax = post
                     } = (preRat ++) <$> balancePostTax 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 :: Rational -> [RawAllocation] -> Either T.Text [BalAllocation]
balancePostTax bal as
  | null as = Left "no allocations to balance"
  | otherwise = case partitionEithers $ fmap hasVal as of
      ([([empty], nonmissing)], bs) ->
        let s = bal - sumAllocations (nonmissing:bs) in
        if s < 0
        then Left "allocations exceed total"
        else Right $ mapAmts (empty { amtValue = s }:) nonmissing : bs
      ([], _) -> Left "need one blank amount to balance"
      _       -> Left "multiple blank amounts present"
  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

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

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

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

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

insertExpense :: MonadIO m => Expense -> MappingT m ()
insertExpense e@Expense { expFrom = from
                      , expTo = to
                      , expCurrency = cur
                      , expBucket = buc
                      , expAmounts = as
                      } = do
  whenHash CTExpense e $ \c -> do
    ts <- concat <$> mapM (timeAmountToTx from to cur) as
    lift $ mapM_ (insertTxBucket (Just buc) c) ts

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

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

insertStatements :: MonadIO m => Config -> MappingT m ()
insertStatements = mapM_ insertStatement . statements

insertStatement :: MonadIO m => Statement -> MappingT m ()
insertStatement (StmtManual m) = insertManual m
insertStatement (StmtImport i) = insertImport i

insertManual :: MonadIO m => Manual -> MappingT m ()
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
    ts <- mapM tx $ expandDatePat bounds dp
    lift $ mapM_ (insertTx c) ts
  where
    tx day = txPair day from to u (dec2Rat v) e

insertImport :: MonadIO m => Import -> MappingT m ()
insertImport i = whenHash CTImport i $ \c -> do
  bounds <- asks kmStatementInterval
  bs <- readImport i
  -- TODO this isn't efficient, the whole file will be read and maybe no
  -- transactions will be desired
  rs <- mapM resolveTx $ filter (inMaybeBounds bounds . txDate) bs
  lift $ mapM_ (insertTx c) rs

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

txPair :: MonadIO m => Day -> AcntID -> AcntID -> T.Text -> Rational -> T.Text
  -> MappingT m 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 :: MonadIO m => BalTx -> MappingT m KeyTx
resolveTx t@Tx { txSplits = ss } = do
  rs <- catMaybes <$> mapM resolveSplit ss
  return $ t { txSplits = rs }

resolveSplit :: MonadIO m => BalSplit -> MappingT m (Maybe 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 $ case (aid, cid, sign) of
    (Just aid', Just cid', Just sign')
      -> Just $ s { sAcnt = aid'
                  , sCurrency = cid'
                  , sValue = v * fromIntegral (sign2Int sign')
                  }
    _ -> Nothing

insertTxBucket :: MonadIO 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 :: MonadIO m => Key CommitR -> KeyTx -> SqlPersistT m ()
insertTx = insertTxBucket Nothing

insertSplit :: MonadIO 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