pwncash/lib/Internal/Database/Ops.hs

{-# LANGUAGE FlexibleContexts  #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards   #-}
{-# LANGUAGE TupleSections     #-}
{-# LANGUAGE TypeApplications  #-}
{-# LANGUAGE TypeFamilies      #-}

module Internal.Database.Ops
  ( migrate_
  , nukeTables
  , showBalances
  , updateHashes
  , getDBState
  , tree2Records
  , flattenAcntRoot
  , paths2IDs

  ) where

import           Control.Monad.Logger
import           Control.Monad.Trans.Reader

import           Conduit

import           Data.Bifunctor
import           Data.Either
import           Data.Hashable
import           Data.List                       ((\\))
import qualified Data.List                       as L
import qualified Data.Map                        as M
import qualified Data.Text                       as T

import           Database.Esqueleto.Experimental
import           Database.Persist.Sql            hiding (delete, (==.), (||.))
import           Database.Persist.Sqlite         hiding (delete, (==.), (||.))
import           Database.Sqlite                 hiding (Config)

import           Internal.Database.Model
import           Internal.Types
import           Internal.Utils

migrate_ :: SqlConfig -> SqlPersistT (ResourceT (NoLoggingT IO)) () -> IO ()
migrate_ c more = runNoLoggingT $ runResourceT
          $ withSqlConn (openConnection c) (\backend ->
                                              flip runSqlConn backend $ do
                                               runMigration migrateAll
                                               more
                                       )

openConnection :: SqlConfig -> LogFunc -> IO SqlBackend
openConnection c logfn = case c of
  Sqlite p -> do
    conn <- open p
    wrapConnection conn logfn
  Postgres -> error "postgres not implemented"

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

showBalances :: MonadIO 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 = Budget { expenses = xs, income = is }
                  , statements = ss } =
  (hash <$> xs) ++ (hash <$> is) ++ (hash <$> ms) ++ (hash <$> ps)
  where
    (ms, ps) = partitionEithers $ fmap go ss
    go (StmtManual x) = Left x
    go (StmtImport 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 :: MonadIO m => SqlPersistT m [Int]
getDBHashes = fmap (commitRHash . entityVal) <$> dumpTbl

nukeDBHash :: MonadIO m => Int -> SqlPersistT m ()
nukeDBHash h = delete $ do
    c <- from table
    where_ (c ^. CommitRHash ==. val h)

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

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

updateHashes :: MonadIO m => Config -> SqlPersistT m [Int]
updateHashes c = do
  (del, new) <- getConfigHashes c
  nukeDBHashes del
  return new

dumpTbl :: (PersistEntity r, MonadIO m) => SqlPersistT m [Entity r]
dumpTbl = select $ from table

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

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

updateAccounts :: MonadIO m => AccountRoot -> SqlPersistT m AccountMap
updateAccounts ar = do
  let (acnts, paths, acntMap) = indexAcntRoot ar
  acnts' <- dumpTbl
  let (toIns, toDel) = setDiff acnts acnts'
  deleteWhere ([] :: [Filter AccountPathR])
  mapM_ deleteAccount toDel
  -- liftIO $ mapM_ print toDel
  mapM_ insertFull toIns
  mapM_ insert paths
  return acntMap

insertFull :: (MonadIO m, PersistStoreWrite b, PersistRecordBackend r b)
  => Entity r -> ReaderT b m ()
insertFull (Entity k v) = insertKey k v

updateCurrencies :: MonadIO m => [Currency] -> SqlPersistT m CurrencyMap
updateCurrencies cs = do
  let curs = fmap currency2Record cs
  curs' <- select $ from $ table @CurrencyR
  let (toIns, toDel) = setDiff curs curs'
  mapM_ deleteCurrency toDel
  mapM_ insertFull toIns
  return $ currencyMap curs

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

currencyMap :: [Entity CurrencyR] -> CurrencyMap
currencyMap = M.fromList . fmap (\e -> (currencyRSymbol $ entityVal e, entityKey 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))])
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) = 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))]
         )
    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
  . unzip
  . L.sortOn fst
  . fmap (first pathList)
  where
    pathList (AcntPath t ns) = reverse $ atName t : ns

trimNames :: [[T.Text]] -> [AcntID]
trimNames = fmap fmt . trimAll 0
  where
    fmt [] = err "blank path"
    fmt ys = T.intercalate "_" $ reverse ys
    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 = take (i + 1)
    err msg = errorWithoutStackTrace $ "Import.Database.Ops.hs: " ++ msg

(!?) :: [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_ {..} =
  ((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) = unzip3 $ uncurry tree2Records <$> flattenAcntRoot r

getDBState :: MonadIO m => Config -> SqlPersistT m (FilePath -> DBState)
getDBState c = do
  am <- updateAccounts $ accounts c
  cm <- updateCurrencies $ currencies c
  hs <- updateHashes c
  -- TODO not sure how I feel about this, probably will change this struct alot
  -- in the future so whatever...for now
  return $ \f -> DBState
    { kmCurrency = cm
    , kmAccount = am
    , kmBudgetInterval = intervalMaybeBounds $ budgetInterval $ global c
    , kmStatementInterval = intervalMaybeBounds $ statementInterval $ global c
    , kmNewCommits = hs
    , kmConfigDir = f
    }