pwncash/lib/Internal/Statement.hs

{-# LANGUAGE RecordWildCards #-}

module Internal.Statement
  ( readImport
  )
where

import Data.Csv
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 :: MonadFinance m => Statement -> m (EitherErrs [BalTx])
readImport Statement {..} = do
  let ores = plural $ compileOptions stmtTxOpts
  let cres = concatEithersL $ compileMatch <$> stmtParsers
  m <- askDBState kmCurrency
  case concatEithers2 ores cres (,) of
    Right (compiledOptions, compiledMatches) -> do
      ires <- mapM (readImport_ stmtSkipLines stmtDelim compiledOptions) stmtPaths
      case concatEitherL ires of
        Right records -> return $ runReader (matchRecords compiledMatches $ L.sort $ concat records) m
        Left es -> return $ Left es
    Left es -> return $ Left es

readImport_
  :: MonadFinance m
  => Natural
  -> Word
  -> TxOptsRe
  -> FilePath
  -> m (EitherErr [TxRecord])
readImport_ n delim tns p = do
  dir <- askDBState kmConfigDir
  bs <- liftIO $ BL.readFile $ dir </> p
  case decodeByNameWithP (parseTxRecord p tns) opts $ skip bs of
    Left m -> return $ Left $ ParseError $ T.pack m
    Right (_, v) -> return $ Right $ 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 :: FilePath -> TxOptsRe -> NamedRecord -> Parser (Maybe TxRecord)
parseTxRecord p 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 p

matchRecords :: [MatchRe] -> [TxRecord] -> CurrencyM (EitherErrs [BalTx])
matchRecords ms rs = do
  res <- matchAll (matchPriorities ms) rs
  case res of
    Left es -> return $ Left es
    Right (matched, unmatched, notfound) -> do
      case (matched, unmatched, notfound) of
        (ms_, [], []) -> do
          -- TODO record number of times each match hits for debugging
          return $ first (: []) $ mapM balanceTx ms_
        (_, us, ns) -> return $ Left [StatementError us ns]

matchPriorities :: [MatchRe] -> [MatchGroup]
matchPriorities =
  fmap matchToGroup
    . L.groupBy (\a b -> spPriority a == spPriority b)
    . L.sortOn (Down . spPriority)

matchToGroup :: [MatchRe] -> MatchGroup
matchToGroup ms =
  uncurry MatchGroup $
    first (L.sortOn spDate) $
      L.partition (isJust . spDate) ms

-- TDOO could use a better struct to flatten the maybe date subtype
data MatchGroup = MatchGroup
  { mgDate :: ![MatchRe]
  , mgNoDate :: ![MatchRe]
  }
  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 MatchRe
  -> TxRecord
  -> CurrencyM (EitherErrs (Zipped MatchRe, MatchRes RawTx))
zipperMatch (Unzipped bs cs as) x = go [] cs
  where
    go _ [] = return $ Right (Zipped bs $ cs ++ as, MatchFail)
    go prev (m : ms) = do
      res <- matches m x
      case res of
        Right MatchFail -> go (m : prev) ms
        Right skipOrPass ->
          let ps = reverse prev
              ms' = maybe ms (: ms) (matchDec m)
           in return $ Right (Zipped bs $ ps ++ ms' ++ as, skipOrPass)
        Left es -> return $ Left es

-- TODO all this unpacking left/error crap is annoying
zipperMatch'
  :: Zipped MatchRe
  -> TxRecord
  -> CurrencyM (EitherErrs (Zipped MatchRe, MatchRes RawTx))
zipperMatch' z x = go z
  where
    go (Zipped bs (a : as)) = do
      res <- matches a x
      case res of
        Right MatchFail -> go (Zipped (a : bs) as)
        Right skipOrPass ->
          return $ Right (Zipped (maybe bs (: bs) $ matchDec a) as, skipOrPass)
        Left es -> return $ Left es
    go z' = return $ Right (z', MatchFail)

matchDec :: MatchRe -> Maybe MatchRe
matchDec m = case spTimes m of
  Just 1 -> Nothing
  Just n -> Just $ m {spTimes = Just $ n - 1}
  Nothing -> Just m

matchAll :: [MatchGroup] -> [TxRecord] -> CurrencyM (EitherErrs ([RawTx], [TxRecord], [MatchRe]))
matchAll = go ([], [])
  where
    go (matched, unused) gs rs = case (gs, rs) of
      (_, []) -> return $ Right (matched, [], unused)
      ([], _) -> return $ Right (matched, rs, unused)
      (g : gs', _) -> do
        res <- matchGroup g rs
        case res of
          Right (ts, unmatched, us) ->
            go (ts ++ matched, us ++ unused) gs' unmatched
          Left es -> return $ Left es

matchGroup :: MatchGroup -> [TxRecord] -> CurrencyM (EitherErrs ([RawTx], [TxRecord], [MatchRe]))
matchGroup MatchGroup {mgDate = ds, mgNoDate = ns} rs = do
  res <- matchDates ds rs
  case res of
    Left es -> return $ Left es
    Right (md, rest, ud) -> do
      res' <- matchNonDates ns rest
      case res' of
        Right (mn, unmatched, un) -> do
          return $ Right $ (md ++ mn, unmatched, filter ((/= Nothing) . spTimes) $ ud ++ un)
        Left es -> return $ Left es

matchDates :: [MatchRe] -> [TxRecord] -> CurrencyM (EitherErrs ([RawTx], [TxRecord], [MatchRe]))
matchDates ms = go ([], [], initZipper ms)
  where
    go (matched, unmatched, z) [] =
      return $
        Right
          ( catMaybes matched
          , reverse unmatched
          , recoverZipper z
          )
    go (matched, unmatched, z) (r : rs) =
      case zipperSlice findDate r z of
        Left zipped -> go (matched, r : unmatched, zipped) rs
        Right unzipped -> do
          res <- zipperMatch unzipped r
          case res of
            Right (z', res') -> do
              let (m, u) = case res' of
                    (MatchPass p) -> (Just p : matched, unmatched)
                    MatchSkip -> (Nothing : matched, unmatched)
                    MatchFail -> (matched, r : unmatched)
              go (m, u, z') rs
            Left es -> return $ Left es
    findDate m r = maybe EQ (`compareDate` trDate r) $ spDate m

matchNonDates :: [MatchRe] -> [TxRecord] -> CurrencyM (EitherErrs ([RawTx], [TxRecord], [MatchRe]))
matchNonDates ms = go ([], [], initZipper ms)
  where
    go (matched, unmatched, z) [] =
      return $
        Right
          ( catMaybes matched
          , reverse unmatched
          , recoverZipper z
          )
    go (matched, unmatched, z) (r : rs) = do
      res <- zipperMatch' z r
      case res of
        Left es -> return $ Left es
        Right (z', res') -> do
          let (m, u) = case res' of
                MatchPass p -> (Just p : matched, unmatched)
                MatchSkip -> (Nothing : matched, unmatched)
                MatchFail -> (matched, r : unmatched)
           in go (m, u, resetZipper z') rs

balanceTx :: RawTx -> EitherErr BalTx
balanceTx t@Tx {txSplits = ss} = do
  bs <- balanceSplits ss
  return $ t {txSplits = bs}

balanceSplits :: [RawSplit] -> EitherErr [BalSplit]
balanceSplits ss =
  fmap concat
    <$> mapM (uncurry bal)
    $ groupByKey
    $ fmap (\s -> (eCurrency s, s)) ss
  where
    haeValue s@Entry {eValue = Just v} = Right s {eValue = v}
    haeValue s = Left s
    bal cur rss
      | length rss < 2 = Left $ BalanceError TooFewSplits cur rss
      | otherwise = case partitionEithers $ fmap haeValue rss of
          ([noVal], val) -> Right $ noVal {eValue = foldr (\s x -> x - eValue s) 0 val} : val
          ([], val) -> Right val
          _ -> Left $ BalanceError NotOneBlank cur rss

groupByKey :: Ord k => [(k, v)] -> [(k, [v])]
groupByKey = M.toList . M.fromListWith (++) . fmap (second (: []))