pwncash/lib/Internal/Statement.hs

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

module Internal.Statement
  ( readImport
  ) where

import           Control.Monad.IO.Class
import           Control.Monad.Trans.Reader

import           Data.Bifunctor
import qualified Data.ByteString.Lazy       as BL
import           Data.Csv
import           Data.Either
import qualified Data.List                  as L
import qualified Data.Map                   as M
import           Data.Maybe
import           Data.Ord
import qualified Data.Text                  as T
import qualified Data.Text.Encoding         as TE
import           Data.Time
import qualified Data.Vector                as V

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

import           Numeric.Natural

import           System.FilePath

-- TODO this probably won't scale well (pipes?)

readImport :: MonadIO 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_ :: MonadIO 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 .: TE.encodeUtf8 toDate
  if d == ""
    then return Nothing
    else do
    a <- parseRational toAmountFmt =<< r .: TE.encodeUtf8 toAmount
    e <- r .: TE.encodeUtf8 toDesc
    os <- M.fromList <$> mapM (\n -> (n, ) <$> r .: TE.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 (:[]))