pwncash/lib/Internal/Insert.hs

{-# LANGUAGE GADTs #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
{-# 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.Map as M

import qualified RIO.List as L
import qualified RIO.Text as T
import RIO.Time

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

-- expandDatePat :: MonadUnliftIO m => Bounds -> DatePat -> MappingT m [Day]
-- expandDatePat d p = do
--   -- TODO crude memoization
--   v <- asks kmBoundsCache
--   modifyMVar v $ \m -> case M.lookup (d, p) m of
--     Just ds -> return (m, ds)
--     Nothing -> do
--       let res = expandDatePat_ d p
--       return (M.insert (d, p) res m, res)

expandDatePat :: Bounds -> DatePat -> [Day]
expandDatePat b (Cron cp) = expandCronPat b cp
-- expandDatePat (a, b) (Cron cp) =
--   fmap xGregToDay $
--     filter (cronPatternMatches cp) $
--       take (fromIntegral $ diffDays b a) $
--         gregorians a
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

-- nextXGreg_ :: CronPat -> XGregorian -> XGregorian
-- nextXGreg_ c XGregorian {xgYear = y, xgMonth = m, xgDay = d, xgDayOfWeek = w}
--   | m == 12 && d == 31 = XGregorian (y + 1) 1 1 w_
--   | (m == 2 && (not leap && d == 28 || (leap && d == 29)))
--       || (m `elem` [4, 6, 9, 11] && d == 30)
--       || (d == 31) =
--       XGregorian y (m + 1) 1 w_
--   | otherwise = XGregorian y m (d + 1) w_
--   where
--     -- don't use DayOfWeek from Data.Time since this uses mod (which uses a
--     -- division opcode) and thus will be slower than just checking for equality
--     -- and adding
--     w_ = if w == 6 then 0 else w + 1
--     leap = isLeapYear $ fromIntegral y

monthLength :: (Integral a, Integral b, Integral c) => a -> b -> c
monthLength y m
  | m == 2 && isLeapYear (fromIntegral y) = 29
  | m == 2 = 28
  | m `elem` [4, 6, 9, 11] = 30
  | otherwise = 31

-- TODO this can be optimized to prevent filtering a bunch of dates for
-- one/a few cron patterns
-- cronPatternMatches :: CronPat -> XGregorian -> Bool
-- cronPatternMatches CronPat {..} XGregorian {..} =
--   testYMD xgYear cronYear
--     && testYMD xgMonth cronMonth
--     && testYMD xgDay cronDay
--     && testW (dayOfWeek_ xgDayOfWeek) cronWeekly
--   where
--     testYMD z = maybe True (mdyPatternMatches (fromIntegral z))
--     testW z = maybe True (`weekdayPatternMatches` z)

expandCronPat :: Bounds -> CronPat -> [Day]
expandCronPat b = L.unfoldr nextCronPat . compileCronPat b

data CompiledCronPat = CompiledCronPat
  { ccpYear :: ![Int]
  , ccpMonth :: !(Zipper Int)
  , ccpDay :: !(Zipper Int)
  , ccpWeekly :: ![Int]
  , ccpMonthEnd :: !Int
  , ccpDayEnd :: !Int
  }
  deriving (Show)

data Zipper a = Zipper ![a] ![a] deriving (Show)

initZipper :: [a] -> Zipper a
initZipper = Zipper []

resetZipper :: Zipper a -> Zipper a
resetZipper (Zipper bs as) = initZipper $ reverse bs ++ as

shiftZipperWhile :: (a -> Bool) -> Zipper a -> Zipper a
shiftZipperWhile f z@(Zipper bs as) = case as of
  [] -> z
  x : xs
    | f x -> shiftZipperWhile f $ Zipper (x : bs) xs
    | otherwise -> z

zipperCurrent :: Zipper a -> Either (Zipper a) (a, Zipper a)
zipperCurrent z@(Zipper _ []) = Left $ resetZipper z
zipperCurrent (Zipper bs (a : as)) = Right (a, Zipper (a : bs) as)

compileCronPat :: Bounds -> CronPat -> CompiledCronPat
compileCronPat (x, y) CronPat {..} =
  CompiledCronPat
    { ccpYear = maybe [y0_ .. y1_] compileMDY_ cronYear
    , ccpMonth = compileDY [1 .. 12] m0 cronMonth
    , ccpDay = compileDY [1 .. 31] d0 cronDay
    , ccpWeekly = maybe [] compileW cronWeekly
    , ccpMonthEnd = m1
    , ccpDayEnd = d1
    }
  where
    (y0, m0, d0) = toGregorian x
    (y1, m1, d1) = toGregorian y
    y0_ = fromIntegral y0
    y1_ = fromIntegral y1
    compileDY def k = shiftZipperWhile (< k) . initZipper . maybe def compileMDY_
    compileMDY_ (Single z) = [fromIntegral z]
    compileMDY_ (Multi zs) = fromIntegral <$> zs
    compileMDY_ (Repeat RepeatPat {rpStart = s, rpBy = b, rpRepeats = rs}) =
      -- TODO minor perf improvement, filter the repeats before filterng <=31
      let b' = fromIntegral b
          xs = takeWhile (<= 31) $ L.iterate (+ b') $ fromIntegral s
       in maybe xs (\r -> take (fromIntegral r) xs) rs

    compileW (OnDay w) = [fromEnum w]
    compileW (OnDays ws) = fromEnum <$> ws

nextCronPat :: CompiledCronPat -> Maybe (Day, CompiledCronPat)
nextCronPat CompiledCronPat {ccpYear = []} = Nothing
nextCronPat c@(CompiledCronPat {..}) =
  case zipperCurrent ccpMonth of
    Left mz -> nextCronPat $ c {ccpYear = ys, ccpMonth = mz, ccpDay = resetZipper ccpDay}
    Right (m, mz) -> case zipperCurrent ccpDay of
      Left dz -> nextCronPat $ c {ccpMonth = mz, ccpDay = dz}
      Right (d, dz)
        | null ys && m >= ccpMonthEnd && d >= ccpDayEnd -> Nothing
        | otherwise -> case dayMaybe m d of
            Nothing -> nextCronPat $ c {ccpMonth = mz, ccpDay = resetZipper dz}
            Just day -> Just (day, c {ccpDay = dz})
  where
    y : ys = ccpYear
    -- TODO not the most efficient way to check weekdays (most likely) since
    -- I have to go through all the trouble of converting to a day and then
    -- doing some complex math to figure out which day of the week it is
    validWeekday day =
      null ccpWeekly
        || (not (null ccpWeekly) && dayToWeekday day `elem` ccpWeekly)
    dayMaybe m d
      | d > monthLength y m = Nothing
      | otherwise =
          let day = fromGregorian (fromIntegral y) m d
           in if validWeekday day then Just day else Nothing

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

-- -- TODO could clean this up by making an enum instance for Weekday
-- dayOfWeek_ :: Int -> Weekday
-- dayOfWeek_ d = case d of
--   0 -> Sun
--   1 -> Mon
--   2 -> Tue
--   3 -> Wed
--   4 -> Thu
--   5 -> Fri
--   _ -> Sat

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

-- mdyPatternMatches :: Natural -> MDYPat -> Bool
-- mdyPatternMatches x p = case p of
--   Single y -> x == y
--   Multi xs -> x `elem` xs
--   Repeat (RepeatPat {rpStart = s, rpBy = b, rpRepeats = r}) ->
--     s >= x && mod x b == 0 && maybe True (\y -> x <= s + b * y) r

withDates
  :: MonadUnliftIO m
  => DatePat
  -> (Day -> MappingT m a)
  -> MappingT m [a]
withDates dp f = do
  bounds <- askBounds
  let days = expandDatePat bounds dp
  mapM f days

askBounds :: MonadUnliftIO m => MappingT m Bounds
askBounds = (liftIO . resolveBounds) =<< asks kmBudgetInterval

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

insertBudget :: MonadUnliftIO m => Budget -> MappingT m [InsertError]
insertBudget Budget {income = is, transfers = es} = do
  es1 <- mapM insertIncome is
  es2 <- mapM insertTransfer 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

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

data BudgetMeta = BudgetMeta
  { bmCommit :: !(Key CommitR)
  , bmWhen :: !Day
  , bmCur :: !CurID
  }

data BudgetTx = BudgetTx
  { btMeta :: !BudgetMeta
  , btFrom :: !(BudgetSplit IncomeBucket)
  , btTo :: !(BudgetSplit ExpenseBucket)
  , btValue :: !Rational
  , btDesc :: !T.Text
  }

insertIncome :: MonadUnliftIO m => Income -> MappingT m [InsertError]
insertIncome i@Income {..} =
  whenHash CTIncome i [] $ \c ->
    unlessLeft (balanceIncome i) $ \balance ->
      fmap concat $ withDates incWhen $ \day -> do
        let meta = BudgetMeta c day incCurrency
        let fromAllos b = concatMap (fromAllo meta incFrom (Just b))
        let pre = fromAllos PreTax incPretax
        let tax = fmap (fromTax meta incFrom) incTaxes
        let post = fromAllos PostTax incPosttax
        let bal =
              BudgetTx
                { btMeta = meta
                , btFrom = BudgetSplit incFrom $ Just PostTax
                , btTo = BudgetSplit incToBal Nothing
                , btValue = balance
                , btDesc = "balance after deductions"
                }
        fmap concat $ mapM insertBudgetTx $ bal : (pre ++ tax ++ post)

fromAllo
  :: BudgetMeta
  -> AcntID
  -> Maybe IncomeBucket
  -> Allocation
  -> [BudgetTx]
fromAllo meta from ib Allocation {..} = fmap (toBT alloPath) alloAmts
  where
    toBT to (Amount desc v) =
      BudgetTx
        { btFrom = BudgetSplit from ib
        , btTo = BudgetSplit to $ Just alloBucket
        , btValue = dec2Rat v
        , btDesc = desc
        , btMeta = meta
        }

fromTax :: BudgetMeta -> AcntID -> Tax -> BudgetTx
fromTax meta from Tax {taxAcnt = to, taxValue = v} =
  BudgetTx
    { btFrom = BudgetSplit from (Just IntraTax)
    , btTo = BudgetSplit to (Just Fixed)
    , btValue = dec2Rat v
    , btDesc = ""
    , btMeta = meta
    }

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)

insertTransfer :: MonadUnliftIO m => Transfer -> MappingT m [InsertError]
insertTransfer t@Transfer {..} =
  fmap (concat . concat) $ whenHash CTExpense t [] $ \key -> do
    forM transAmounts $ \(TimeAmount amt pat) ->
      withDates pat $ \day -> insertBudgetTx $ budgetTx amt day key
  where
    meta d c = BudgetMeta {bmWhen = d, bmCur = transCurrency, bmCommit = c}
    budgetTx (Amount desc v) d c =
      BudgetTx
        { btMeta = meta d c
        , btFrom = BudgetSplit transFrom Nothing
        , btTo = BudgetSplit transTo Nothing
        , btValue = dec2Rat v
        , btDesc = desc
        }

insertBudgetTx :: MonadUnliftIO m => BudgetTx -> MappingT m [InsertError]
insertBudgetTx BudgetTx {..} = do
  res <- splitPair (bsAcnt btFrom) (bsAcnt btTo) (bmCur btMeta) btValue
  unlessLefts_ res $ \(sFrom, sTo) -> lift $ do
    k <- insert $ TransactionR (bmCommit btMeta) (bmWhen btMeta) btDesc
    skFrom <- insertSplit k sFrom
    bFrom <- insert $ BudgetLabelR skFrom ""
    forM_ (bsBucket btFrom) $ \b ->
      insert_ $ IncomeBucketR bFrom b
    skTo <- insertSplit k sTo
    bTo <- insert $ BudgetLabelR skTo ""
    forM_ (bsBucket btTo) $ \b ->
      insert_ $ ExpenseBucketR bTo b

splitPair
  :: MonadUnliftIO m
  => AcntID
  -> AcntID
  -> CurID
  -> Rational
  -> MappingT 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
          }

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

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

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
      let days = expandDatePat bounds dp
      res <- mapM tx days
      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
  -> CurID
  -> 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_)
      }

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