pwncash/lib/Internal/Database.hs

module Internal.Database
  ( runDB
  , nukeTables
  , updateHashes
  , updateDBState
  , getDBState
  , tree2Records
  , flattenAcntRoot
  , paths2IDs
  , mkPool
  , whenHash0
  , whenHash
  , whenHash_
  , eitherHash
  , insertEntry
  , resolveEntry
  , readUpdates
  )
where

import Conduit
import Control.Monad.Except
import Control.Monad.Logger
import Data.Hashable
import Database.Esqueleto.Experimental ((:&) (..), (==.), (^.))
import qualified Database.Esqueleto.Experimental as E
import Database.Esqueleto.Internal.Internal (SqlSelect)
import Database.Persist.Monad
import Database.Persist.Sqlite hiding
  ( delete
  , deleteWhere
  , insert
  , insertKey
  , insert_
  , runMigration
  , (==.)
  , (||.)
  )
import GHC.Err
import Internal.Types.Main
import Internal.Utils
import RIO hiding (LogFunc, isNothing, on, (^.))
import RIO.List ((\\))
import qualified RIO.List as L
import qualified RIO.Map as M
import qualified RIO.NonEmpty as N
import qualified RIO.Text as T
import qualified RIO.Vector as V

runDB
  :: MonadUnliftIO m
  => SqlConfig
  -> SqlQueryT (NoLoggingT m) a
  -> m a
runDB c more =
  runNoLoggingT $ do
    pool <- mkPool c
    runSqlQueryT pool $ do
      _ <- lift askLoggerIO
      runMigration migrateAll
      more

mkPool :: (MonadLoggerIO m, MonadUnliftIO m) => SqlConfig -> m ConnectionPool
mkPool c = case c of
  Sqlite p -> createSqlitePool p 10
  -- conn <- open p
  -- wrapConnection conn logfn
  Postgres -> error "postgres not implemented"

nukeTables :: MonadSqlQuery m => m ()
nukeTables = do
  deleteWhere ([] :: [Filter CommitR])
  deleteWhere ([] :: [Filter CurrencyR])
  deleteWhere ([] :: [Filter AccountR])
  deleteWhere ([] :: [Filter TransactionR])

-- showBalances :: MonadUnliftIO m => SqlPersistT m ()
-- showBalances = do
--   xs <- select $ do
--     (accounts :& splits :& txs) <-
--       from
--         $ table @AccountR
--           `innerJoin` table @SplitR
--         `on` (\(a :& s) -> a ^. AccountRId ==. s ^. SplitRAccount)
--           `innerJoin` table @TransactionR
--         `on` (\(_ :& s :& t) -> s ^. SplitRTransaction ==. t ^. TransactionRId)
--     where_ $
--       isNothing (txs ^. TransactionRBucket)
--         &&. ( (accounts ^. AccountRFullpath `like` val "asset" ++. (%))
--                 ||. (accounts ^. AccountRFullpath `like` val "liability" ++. (%))
--             )
--     groupBy (accounts ^. AccountRFullpath, accounts ^. AccountRName)
--     return
--       ( accounts ^. AccountRFullpath
--       , accounts ^. AccountRName
--       , sum_ $ splits ^. SplitRValue
--       )
--   -- TODO super stetchy table printing thingy
--   liftIO $ do
--     putStrLn $ T.unpack $ fmt "Account" "Balance"
--     putStrLn $ T.unpack $ fmt (T.replicate 60 "-") (T.replicate 15 "-")
--     mapM_ (putStrLn . T.unpack . fmtBalance) xs
--   where
--     fmtBalance (path, name, bal) = fmt (toFullPath path name) (toBal bal)
--     fmt a b = T.unwords ["| ", pad 60 a, " | ", pad 15 b, " |"]
--     pad n xs = T.append xs $ T.replicate (n - T.length xs) " "
--     toFullPath path name = T.unwords [unValue @T.Text path, "/", unValue @T.Text name]
--     toBal = maybe "???" (fmtRational 2) . unValue

hashConfig :: Config -> [Int]
hashConfig
  Config_
    { budget = bs
    , statements = ss
    } = (hash <$> bs) ++ (hash <$> ms) ++ (hash <$> ps)
    where
      (ms, ps) = partitionEithers $ fmap go ss
      go (HistTransfer x) = Left x
      go (HistStatement x) = Right x

setDiff :: Eq a => [a] -> [a] -> ([a], [a])
-- setDiff = setDiff' (==)
setDiff as bs = (as \\ bs, bs \\ as)

-- setDiff' :: Eq a => (a -> b -> Bool) -> [a] -> [b] -> ([a], [b])
-- setDiff' f = go []
--   where
--     go inA [] bs     = (inA, bs)
--     go inA as []     = (as ++ inA, [])
--     go inA (a:as) bs = case inB a bs of
--       Just bs' -> go inA as bs'
--       Nothing  -> go (a:inA) as bs
--     inB _ [] = Nothing
--     inB a (b:bs)
--       | f a b = Just bs
--       | otherwise = inB a bs

getDBHashes :: MonadSqlQuery m => m [Int]
getDBHashes = fmap (commitRHash . entityVal) <$> dumpTbl

nukeDBHash :: MonadSqlQuery m => Int -> m ()
nukeDBHash h = deleteE $ do
  c <- E.from E.table
  E.where_ (c ^. CommitRHash ==. E.val h)

nukeDBHashes :: MonadSqlQuery m => [Int] -> m ()
nukeDBHashes = mapM_ nukeDBHash

getConfigHashes :: MonadSqlQuery m => Config -> m ([Int], [Int])
getConfigHashes c = do
  let ch = hashConfig c
  dh <- getDBHashes
  return $ setDiff dh ch

dumpTbl :: (MonadSqlQuery m, PersistEntity r) => m [Entity r]
dumpTbl = selectE $ E.from E.table

deleteAccount :: MonadSqlQuery m => Entity AccountR -> m ()
deleteAccount e = deleteE $ do
  c <- E.from $ E.table @AccountR
  E.where_ (c ^. AccountRId ==. E.val k)
  where
    k = entityKey e

deleteCurrency :: MonadSqlQuery m => Entity CurrencyR -> m ()
deleteCurrency e = deleteE $ do
  c <- E.from $ E.table @CurrencyR
  E.where_ (c ^. CurrencyRId ==. E.val k)
  where
    k = entityKey e

deleteTag :: MonadSqlQuery m => Entity TagR -> m ()
deleteTag e = deleteE $ do
  c <- E.from $ E.table @TagR
  E.where_ (c ^. TagRId ==. E.val k)
  where
    k = entityKey e

-- TODO slip-n-slide code...
insertFull
  :: (PersistRecordBackend r SqlBackend, Typeable r, MonadSqlQuery m)
  => Entity r
  -> m ()
insertFull (Entity k v) = insertKey k v

currency2Record :: Currency -> Entity CurrencyR
currency2Record c@Currency {curSymbol, curFullname, curPrecision} =
  Entity (toKey c) $ CurrencyR curSymbol curFullname (fromIntegral curPrecision)

currencyMap :: [Entity CurrencyR] -> CurrencyMap
currencyMap =
  M.fromList
    . fmap
      ( \e ->
          ( currencyRSymbol $ entityVal e
          , (entityKey e, fromIntegral $ currencyRPrecision $ entityVal e)
          )
      )

toKey :: (ToBackendKey SqlBackend b, Hashable a) => a -> Key b
toKey = toSqlKey . fromIntegral . hash

tree2Entity :: AcntType -> [T.Text] -> T.Text -> T.Text -> Entity AccountR
tree2Entity t parents name des =
  Entity (toSqlKey $ fromIntegral h) $
    AccountR name (toPath parents) des
  where
    p = AcntPath t (reverse (name : parents))
    h = hash p
    toPath = T.intercalate "/" . (atName t :) . reverse

tree2Records
  :: AcntType
  -> AccountTree
  -> ([Entity AccountR], [AccountPathR], [(AcntPath, (AccountRId, AcntSign, AcntType))])
tree2Records t = go []
  where
    go ps (Placeholder d n cs) =
      let e = tree2Entity t (fmap snd ps) n d
          k = entityKey e
          (as, aps, ms) = L.unzip3 $ fmap (go ((k, n) : ps)) cs
          a0 = acnt k n (fmap snd ps) d
          paths = expand k $ fmap fst ps
       in (a0 : concat as, paths ++ concat aps, concat ms)
    go ps (Account d n) =
      let e = tree2Entity t (fmap snd ps) n d
          k = entityKey e
       in ( [acnt k n (fmap snd ps) d]
          , expand k $ fmap fst ps
          , [(AcntPath t $ reverse $ n : fmap snd ps, (k, sign, t))]
          )
    toPath = T.intercalate "/" . (atName t :) . reverse
    acnt k n ps = Entity k . AccountR n (toPath ps)
    expand h0 hs = (\(h, d) -> AccountPathR h h0 d) <$> zip (h0 : hs) [0 ..]
    sign = accountSign t

paths2IDs :: [(AcntPath, a)] -> [(AcntID, a)]
paths2IDs =
  uncurry zip
    . first trimNames
    . L.unzip
    . L.sortOn fst
    . fmap (first pathList)
  where
    pathList (AcntPath t []) = atName t :| []
    pathList (AcntPath t ns) = N.reverse $ atName t :| ns

-- none of these errors should fire assuming that input is sorted and unique
trimNames :: [N.NonEmpty T.Text] -> [AcntID]
trimNames = fmap (T.intercalate "_" . reverse) . trimAll 0
  where
    trimAll _ [] = []
    trimAll i (y : ys) = case L.foldl' (matchPre i) (y, [], []) ys of
      (a, [], bs) -> reverse $ trim i a : bs
      (a, as, bs) -> reverse bs ++ trimAll (i + 1) (reverse $ a : as)
    matchPre i (y, ys, old) new = case (y !? i, new !? i) of
      (Nothing, Just _) ->
        case ys of
          [] -> (new, [], trim i y : old)
          _ -> err "unsorted input"
      (Just _, Nothing) -> err "unsorted input"
      (Nothing, Nothing) -> err "duplicated inputs"
      (Just a, Just b)
        | a == b -> (new, y : ys, old)
        | otherwise ->
            let next = case ys of
                  [] -> [trim i y]
                  _ -> trimAll (i + 1) (reverse $ y : ys)
             in (new, [], reverse next ++ old)
    trim i = N.take (i + 1)
    err msg = errorWithoutStackTrace $ "Import.Database.Ops.hs: " ++ msg

(!?) :: N.NonEmpty a -> Int -> Maybe a
xs !? n
  | n < 0 = Nothing
  -- Definition adapted from GHC.List
  | otherwise =
      foldr
        ( \x r k -> case k of
            0 -> Just x
            _ -> r (k - 1)
        )
        (const Nothing)
        xs
        n

flattenAcntRoot :: AccountRoot -> [(AcntType, AccountTree)]
flattenAcntRoot AccountRoot_ {arIncome, arExpenses, arLiabilities, arAssets, arEquity} =
  ((IncomeT,) <$> arIncome)
    ++ ((ExpenseT,) <$> arExpenses)
    ++ ((LiabilityT,) <$> arLiabilities)
    ++ ((AssetT,) <$> arAssets)
    ++ ((EquityT,) <$> arEquity)

indexAcntRoot :: AccountRoot -> ([Entity AccountR], [AccountPathR], AccountMap)
indexAcntRoot r =
  ( concat ars
  , concat aprs
  , M.fromList $ paths2IDs $ concat ms
  )
  where
    (ars, aprs, ms) = L.unzip3 $ uncurry tree2Records <$> flattenAcntRoot r

getDBState
  :: (MonadInsertError m, MonadSqlQuery m)
  => Config
  -> m (DBState, DBUpdates)
getDBState c = do
  (del, new) <- getConfigHashes c
  combineError bi si $ \b s ->
    ( DBState
        { kmCurrency = currencyMap cs
        , kmAccount = am
        , kmBudgetInterval = b
        , kmStatementInterval = s
        , kmTag = tagMap ts
        , kmNewCommits = new
        }
    , DBUpdates
        { duOldCommits = del
        , duNewTagIds = ts
        , duNewAcntPaths = paths
        , duNewAcntIds = acnts
        , duNewCurrencyIds = cs
        }
    )
  where
    bi = liftExcept $ resolveDaySpan $ budgetInterval $ global c
    si = liftExcept $ resolveDaySpan $ statementInterval $ global c
    (acnts, paths, am) = indexAcntRoot $ accounts c
    cs = currency2Record <$> currencies c
    ts = toRecord <$> tags c
    toRecord t@Tag {tagID, tagDesc} = Entity (toKey t) $ TagR tagID tagDesc
    tagMap = M.fromList . fmap (\e -> (tagRSymbol $ entityVal e, entityKey e))

updateHashes :: (MonadSqlQuery m) => DBUpdates -> m ()
updateHashes DBUpdates {duOldCommits} = nukeDBHashes duOldCommits

updateTags :: (MonadFinance m, MonadSqlQuery m) => DBUpdates -> m ()
updateTags DBUpdates {duNewTagIds} = do
  tags' <- selectE $ E.from $ E.table @TagR
  let (toIns, toDel) = setDiff duNewTagIds tags'
  mapM_ deleteTag toDel
  mapM_ insertFull toIns

updateAccounts :: (MonadFinance m, MonadSqlQuery m) => DBUpdates -> m ()
updateAccounts DBUpdates {duNewAcntIds, duNewAcntPaths} = do
  acnts' <- dumpTbl
  let (toIns, toDel) = setDiff duNewAcntIds acnts'
  deleteWhere ([] :: [Filter AccountPathR])
  mapM_ deleteAccount toDel
  mapM_ insertFull toIns
  mapM_ insert duNewAcntPaths

updateCurrencies :: (MonadFinance m, MonadSqlQuery m) => DBUpdates -> m ()
updateCurrencies DBUpdates {duNewCurrencyIds} = do
  curs' <- selectE $ E.from $ E.table @CurrencyR
  let (toIns, toDel) = setDiff duNewCurrencyIds curs'
  mapM_ deleteCurrency toDel
  mapM_ insertFull toIns

updateDBState :: (MonadFinance m, MonadSqlQuery m) => DBUpdates -> m ()
updateDBState u = do
  updateHashes u
  updateTags u
  updateAccounts u
  updateCurrencies u

deleteE :: (MonadSqlQuery m) => E.SqlQuery () -> m ()
deleteE q = unsafeLiftSql "esqueleto-delete" (E.delete q)

selectE :: (MonadSqlQuery m, SqlSelect a r) => E.SqlQuery a -> m [r]
selectE q = unsafeLiftSql "esqueleto-select" (E.select q)

whenHash
  :: (Hashable a, MonadFinance m, MonadSqlQuery m)
  => ConfigType
  -> a
  -> b
  -> (CommitRId -> m b)
  -> m b
whenHash t o def f = do
  let h = hash o
  hs <- askDBState kmNewCommits
  if h `elem` hs then f =<< insert (CommitR h t) else return def

whenHash0
  :: (Hashable a, MonadFinance m)
  => ConfigType
  -> a
  -> b
  -> (CommitR -> m b)
  -> m b
whenHash0 t o def f = do
  let h = hash o
  hs <- askDBState kmNewCommits
  if h `elem` hs then f (CommitR h t) else return def

eitherHash
  :: (Hashable a, MonadFinance m)
  => ConfigType
  -> a
  -> (CommitR -> m b)
  -> (CommitR -> m c)
  -> m (Either b c)
eitherHash t o f g = do
  let h = hash o
  let c = CommitR h t
  hs <- askDBState kmNewCommits
  if h `elem` hs then Left <$> f c else Right <$> g c

whenHash_
  :: (Hashable a, MonadFinance m)
  => ConfigType
  -> a
  -> m b
  -> m (Maybe (CommitR, b))
whenHash_ t o f = do
  let h = hash o
  let c = CommitR h t
  hs <- askDBState kmNewCommits
  if h `elem` hs then Just . (c,) <$> f else return Nothing

insertEntry :: MonadSqlQuery m => TransactionRId -> KeyEntry -> m EntryRId
insertEntry
  t
  FullEntry
    { feEntry = Entry {eValue, eTags, eAcnt, eComment}
    , feCurrency
    , feIndex
    , feDeferred
    } =
    do
      k <- insert $ EntryR t feCurrency eAcnt eComment eValue feIndex defval deflink
      mapM_ (insert_ . TagRelationR k) eTags
      return k
    where
      (defval, deflink) = case feDeferred of
        (Just (EntryLinked index scale)) -> (Just scale, Just $ fromIntegral index)
        (Just (EntryBalance target)) -> (Just target, Nothing)
        Nothing -> (Nothing, Nothing)

resolveEntry :: (MonadInsertError m, MonadFinance m) => BalEntry -> m KeyEntry
resolveEntry s@FullEntry {feEntry = e@Entry {eValue, eTags, eAcnt}, feCurrency} = do
  let aRes = lookupAccountKey eAcnt
  let cRes = lookupCurrencyKey feCurrency
  let sRes = lookupAccountSign eAcnt
  let tagRes = combineErrors $ fmap lookupTag eTags
  -- TODO correct sign here?
  -- TODO lenses would be nice here
  combineError (combineError3 aRes cRes sRes (,,)) tagRes $
    \(aid, cid, sign) tags ->
      s
        { feCurrency = cid
        , feEntry = e {eAcnt = aid, eValue = fromIntegral (sign2Int sign) * eValue, eTags = tags}
        }

readUpdates
  :: (MonadInsertError m, MonadSqlQuery m)
  => [Int]
  -> m [Either ReadEntry UpdateEntrySet]
readUpdates hashes = do
  xs <- selectE $ do
    (commits :& txs :& entries) <-
      E.from
        $ E.table @CommitR
          `E.innerJoin` E.table @TransactionR
        `E.on` (\(c :& t) -> c ^. CommitRId ==. t ^. TransactionRCommit)
          `E.innerJoin` E.table @EntryR
        `E.on` (\(_ :& t :& e) -> t ^. TransactionRId ==. e ^. EntryRTransaction)
    E.where_ $ commits ^. CommitRHash `E.in_` E.valList hashes
    return
      ( txs ^. TransactionRDeferred
      , txs ^. TransactionRDate
      , entries
      )
  let (toUpdate, toRead) =
        bimap unpack (fmap makeRE . unpack) $
          L.partition (\(d, _, _) -> E.unValue d) xs
  toUpdate' <-
    liftExcept $
      mapErrors makeUES $
        second (fmap snd) <$> groupWith uGroup toUpdate
  return $ fmap Left toRead ++ fmap Right toUpdate'
  where
    unpack = fmap (\(_, d, e) -> (E.unValue d, (entityKey e, entityVal e)))
    uGroup (day, (_, e)) = (day, entryRCurrency e, entryRTransaction e)
    makeUES ((day, cur, _), es) = do
      let (froms, tos) =
            L.partition ((< 0) . entryRIndex . snd) $
              L.sortOn (entryRIndex . snd) es
      let tot = sum $ fmap (entryRValue . snd) froms
      (from0, fromRO, fromUnk, fromVec) <- splitFrom $ reverse froms
      (to0, toRO, toUnk, toLink0, toLinkN) <- splitTo fromVec tos
      return
        UpdateEntrySet
          { utDate = day
          , utCurrency = cur
          , utFrom0 = from0
          , utTo0 = to0
          , utFromRO = fromRO
          , utToRO = toRO
          , utToUnkLink0 = toLink0
          , utPairs = toLinkN
          , utFromUnk = fromUnk
          , utToUnk = toUnk
          , utTotalValue = tot
          }
    makeRE (d, (_, e)) =
      ReadEntry
        { reDate = d
        , reCurrency = entryRCurrency e
        , reAcnt = entryRAccount e
        , reValue = entryRValue e
        }

splitFrom
  :: [(EntryRId, EntryR)]
  -> InsertExcept (UEBlank, [UE_RO], [UEBalance], Vector (Maybe UEBalance))
splitFrom from = do
  -- ASSUME entries are sorted by index
  (primary, rest) <- case from of
    ((i, e) : xs) -> return (makeUnkUE i e, xs)
    _ -> throwError $ InsertException undefined
  let rest' = fmap splitDeferredValue rest
  let idxVec = V.fromList $ fmap (either (const Nothing) Just) rest'
  let (ro, toBal) = partitionEithers rest'
  return (primary, ro, toBal, idxVec)

splitTo
  :: Vector (Maybe UEBalance)
  -> [(EntryRId, EntryR)]
  -> InsertExcept
      ( UEBlank
      , [UE_RO]
      , [UEBalance]
      , [UELink]
      , [(UEBalance, [UELink])]
      )
splitTo froms tos = do
  -- How to split the credit side of the database transaction in 1024 easy
  -- steps:
  --
  -- 1. ASSUME the entries are sorted by index. Isolate the first as the
  --    primary and puke in user's face if list is empty (which it should never
  --    be)
  (primary, rest) <- case tos of
    ((i, e) : xs) -> return (makeUnkUE i e, xs)
    _ -> throwError $ InsertException undefined

  -- 1. Split the entries based on if they have a link
  let (unlinked, linked) = partitionEithers $ fmap splitLinked rest

  -- 2. Split unlinked based on if they have a balance target
  let (ro, toBal) = partitionEithers $ fmap splitDeferredValue unlinked

  -- 3. Split paired entries by link == 0 (which are special) or link > 0
  let (paired0, pairedN) = second (groupKey id) $ L.partition ((== 0) . fst) linked
  let paired0Res = mapErrors (makeLinkUnk . snd) paired0

  -- 4. Group linked entries (which now have links > 0) according to the debit
  --    entry to which they are linked. If the debit entry cannot be found or
  --    if the linked entry has no scale, blow up in user's face. If the
  --    debit entry is read-only (signified by Nothing in the 'from' array)
  --    then consider the linked entry as another credit read-only entry
  let pairedRes = partitionEithers <$> mapErrors splitPaired pairedN

  combineError paired0Res pairedRes $ \paired0' (pairedUnk, pairedRO) ->
    (primary, ro ++ concat pairedRO, toBal, paired0', pairedUnk)
  where
    splitLinked t@(_, e) = maybe (Left t) (Right . (,t)) $ entryRDeferred_link e
    splitPaired (lnk, ts) = case froms V.!? (lnk - 1) of
      Just (Just f) -> Left . (f,) <$> mapErrors makeLinkUnk ts
      Just Nothing -> return $ Right $ makeRoUE . snd <$> ts
      Nothing -> throwError $ InsertException undefined
    makeLinkUnk (k, e) =
      maybe
        (throwError $ InsertException undefined)
        (return . makeUE k e . LinkScale)
        $ entryRDeferred_value e

splitDeferredValue :: (EntryRId, EntryR) -> Either UE_RO UEBalance
splitDeferredValue (k, e) =
  maybe (Left $ makeRoUE e) (Right . fmap BalanceTarget . makeUE k e) $
    entryRDeferred_value e

makeUE :: i -> EntryR -> v -> UpdateEntry i v
makeUE k e v = UpdateEntry k (entryRAccount e) v (entryRIndex e)

makeRoUE :: EntryR -> UpdateEntry () EntryValue
makeRoUE e = makeUE () e $ EntryValue (entryRValue e)

makeUnkUE :: EntryRId -> EntryR -> UpdateEntry EntryRId ()
makeUnkUE k e = makeUE k e ()