pwncash/lib/Internal/Statement.hs

222 lines
7.1 KiB
Haskell

{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TupleSections #-}
module Internal.Statement
( readImport
)
where
import Data.Csv
import Internal.Database.Model
import Internal.Types
import Internal.Utils
import RIO
import qualified RIO.ByteString.Lazy as BL
import RIO.FilePath
import qualified RIO.List as L
import qualified RIO.Map as M
import qualified RIO.Text as T
import RIO.Time
import qualified RIO.Vector as V
-- TODO this probably won't scale well (pipes?)
readImport :: MonadUnliftIO m => Import -> MappingT m [BalTx]
readImport
Import
{ impPaths = ps
, impMatches = ms
, impTxOpts = ns
, impDelim = d
, impSkipLines = n
} = do
rs <- L.sort . concat <$> mapM (readImport_ n d ns) ps
let (ts, es, notfound) = matchRecords ms rs
liftIO $ mapM_ putStrLn $ reverse es
liftIO $ mapM_ print notfound
return ts
readImport_
:: MonadUnliftIO m
=> Natural
-> Word
-> TxOpts
-> FilePath
-> MappingT m [TxRecord]
readImport_ n delim tns p = do
dir <- asks kmConfigDir
bs <- liftIO $ BL.readFile $ dir </> p
case decodeByNameWithP (parseTxRecord tns) opts $ skip bs of
Left m -> liftIO $ putStrLn m >> return []
Right (_, v) -> return $ catMaybes $ V.toList v
where
opts = defaultDecodeOptions {decDelimiter = fromIntegral delim}
skip = BL.intercalate "\n" . L.drop (fromIntegral n) . BL.split 10
-- TODO handle this better, this maybe thing is a hack to skip lines with
-- blank dates but will likely want to make this more flexible
parseTxRecord :: TxOpts -> NamedRecord -> Parser (Maybe TxRecord)
parseTxRecord TxOpts {..} r = do
d <- r .: T.encodeUtf8 toDate
if d == ""
then return Nothing
else do
a <- parseRational toAmountFmt =<< r .: T.encodeUtf8 toAmount
e <- r .: T.encodeUtf8 toDesc
os <- M.fromList <$> mapM (\n -> (n,) <$> r .: T.encodeUtf8 n) toOther
d' <- parseTimeM True defaultTimeLocale (T.unpack toDateFmt) d
return $ Just $ TxRecord d' a e os
matchRecords :: [Match] -> [TxRecord] -> ([BalTx], [String], [Match])
matchRecords ms rs =
( catMaybes ts
, T.unpack <$> (es ++ bu)
, -- TODO record number of times each match hits for debugging
notfound
)
where
(matched, unmatched, notfound) = matchAll (matchPriorities ms) rs
(es, ts) =
partitionEithers $
fmap Just . balanceTx <$> catMaybes matched
bu = fmap (\x -> T.pack $ "unmatched: " ++ show x) unmatched
matchPriorities :: [Match] -> [MatchGroup]
matchPriorities =
fmap matchToGroup
. L.groupBy (\a b -> mPriority a == mPriority b)
. L.sortOn (Down . mPriority)
matchToGroup :: [Match] -> MatchGroup
matchToGroup ms =
uncurry MatchGroup $
first (L.sortOn mDate) $
L.partition (isJust . mDate) ms
-- TDOO could use a better struct to flatten the maybe date subtype
data MatchGroup = MatchGroup
{ mgDate :: [Match]
, mgNoDate :: [Match]
}
deriving (Show)
data Zipped a = Zipped ![a] ![a]
data Unzipped a = Unzipped ![a] ![a] ![a]
initZipper :: [a] -> Zipped a
initZipper = Zipped []
resetZipper :: Zipped a -> Zipped a
resetZipper = initZipper . recoverZipper
recoverZipper :: Zipped a -> [a]
recoverZipper (Zipped as bs) = reverse as ++ bs
zipperSlice :: (a -> b -> Ordering) -> b -> Zipped a -> Either (Zipped a) (Unzipped a)
zipperSlice f x = go
where
go z@(Zipped _ []) = Left z
go z@(Zipped bs (a : as)) = case f a x of
GT -> go $ Zipped (a : bs) as
EQ -> Right $ goEq (Unzipped bs [a] as)
LT -> Left z
goEq z@(Unzipped _ _ []) = z
goEq z@(Unzipped bs cs (a : as)) = case f a x of
GT -> goEq $ Unzipped (a : bs) cs as
EQ -> goEq $ Unzipped bs (a : cs) as
LT -> z
zipperMatch :: Unzipped Match -> TxRecord -> (Zipped Match, Maybe (Maybe RawTx))
zipperMatch (Unzipped bs cs as) x = go [] cs
where
go _ [] = (Zipped bs $ cs ++ as, Nothing)
go prev (m : ms) = case matches m x of
Nothing -> go (m : prev) ms
res@(Just _) ->
let ps = reverse prev
ms' = maybe ms (: ms) (matchDec m)
in (Zipped bs $ ps ++ ms' ++ as, res)
zipperMatch' :: Zipped Match -> TxRecord -> (Zipped Match, Maybe (Maybe RawTx))
zipperMatch' z x = go z
where
go (Zipped bs (a : as)) = case matches a x of
Nothing -> go (Zipped (a : bs) as)
res -> (Zipped (maybe bs (: bs) $ matchDec a) as, res)
go z' = (z', Nothing)
matchDec :: Match -> Maybe Match
matchDec m@Match {mTimes = t} =
if t' == Just 0 then Nothing else Just $ m {mTimes = t'}
where
t' = fmap pred t
matchAll :: [MatchGroup] -> [TxRecord] -> ([Maybe RawTx], [TxRecord], [Match])
matchAll = go ([], [])
where
go (matched, unused) gs rs = case (gs, rs) of
(_, []) -> (matched, [], unused)
([], _) -> (matched, rs, unused)
(g : gs', _) ->
let (ts, unmatched, us) = matchGroup g rs
in go (ts ++ matched, us ++ unused) gs' unmatched
matchGroup :: MatchGroup -> [TxRecord] -> ([Maybe RawTx], [TxRecord], [Match])
matchGroup MatchGroup {mgDate = ds, mgNoDate = ns} rs =
(md ++ mn, unmatched, filter ((/= Nothing) . mTimes) $ ud ++ un)
where
(md, rest, ud) = matchDates ds rs
(mn, unmatched, un) = matchNonDates ns rest
matchDates :: [Match] -> [TxRecord] -> ([Maybe RawTx], [TxRecord], [Match])
matchDates ms = go ([], [], initZipper ms)
where
go (matched, unmatched, z) [] = (matched, reverse unmatched, recoverZipper z)
go (matched, unmatched, z) (r : rs) = case zipperSlice findDate r z of
Left res -> go (matched, r : unmatched, res) rs
Right res ->
let (z', p) = zipperMatch res r
(m, u) = case p of
Just p' -> (p' : matched, unmatched)
Nothing -> (matched, r : unmatched)
in go (m, u, z') rs
findDate m r = maybe EQ (`compareDate` trDate r) $ mDate m
matchNonDates :: [Match] -> [TxRecord] -> ([Maybe RawTx], [TxRecord], [Match])
matchNonDates ms = go ([], [], initZipper ms)
where
go (matched, unmatched, z) [] = (matched, reverse unmatched, recoverZipper z)
go (matched, unmatched, z) (r : rs) =
let (z', res) = zipperMatch' z r
(m, u) = case res of
Just x -> (x : matched, unmatched)
Nothing -> (matched, r : unmatched)
in go (m, u, resetZipper z') rs
balanceTx :: RawTx -> Either T.Text BalTx
balanceTx t@Tx {txSplits = ss} = do
bs <- balanceSplits ss
return $ t {txSplits = bs}
balanceSplits :: [RawSplit] -> Either T.Text [BalSplit]
balanceSplits ss =
fmap concat
<$> mapM (uncurry bal)
$ groupByKey
$ fmap (\s -> (sCurrency s, s)) ss
where
hasValue s@(Split {sValue = Just v}) = Right s {sValue = v}
hasValue s = Left s
bal cur rss
| length rss < 2 = Left $ T.append "Need at least two splits to balance: " cur
| otherwise = case partitionEithers $ fmap hasValue rss of
([noVal], val) -> Right $ noVal {sValue = foldr (\s x -> x - sValue s) 0 val} : val
([], val) -> Right val
_ -> Left $ T.append "Exactly one split must be blank: " cur
groupByKey :: Ord k => [(k, v)] -> [(k, [v])]
groupByKey = M.toList . M.fromListWith (++) . fmap (second (: []))